grpc_client.cc 8.0 KB
Newer Older
1
/* Copyright (c) 2016 PaddlePaddle Authors. All Rights Reserved.
G
gongweibao 已提交
2 3 4 5 6 7 8 9 10 11 12 13 14

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. */

Y
Yi Wang 已提交
15 16
#include "paddle/fluid/operators/detail/grpc_client.h"

17
#include <sys/time.h>
Y
Yi Wang 已提交
18 19 20

#include <limits>

Y
Yi Wang 已提交
21
#include "paddle/fluid/framework/threadpool.h"
X
Xin Pan 已提交
22
#include "paddle/fluid/platform/profiler.h"
23

G
gongweibao 已提交
24 25 26 27 28 29 30 31 32
namespace paddle {
namespace operators {
namespace detail {

bool RPCClient::AsyncSendVariable(const std::string& ep,
                                  const platform::DeviceContext& ctx,
                                  const framework::Scope& scope,
                                  const std::string& var_name,
                                  int64_t time_out) {
33 34 35 36
  const platform::DeviceContext* p_ctx = &ctx;
  const std::string ep_val = ep;
  const std::string var_name_val = var_name;
  const framework::Scope* p_scope = &scope;
Y
Yancey1989 已提交
37
  const auto ch = GetChannel(ep_val);
38

T
wip  
typhoonzero 已提交
39 40
  framework::AsyncIO([var_name_val, p_ctx, ep_val, p_scope, time_out, ch,
                      this] {
41
    auto* var = p_scope->FindVar(var_name_val);
42 43 44

    ::grpc::ByteBuffer req;
    SerializeToByteBuffer(var_name_val, var, *p_ctx, &req);
45 46 47 48 49 50 51 52 53 54 55

    // varhandle
    VarHandle var_h;
    var_h.ep = ep_val;
    var_h.scope = p_scope;
    var_h.name = var_name_val;
    var_h.ctx = p_ctx;

    // stub context
    SendProcessor* s = new SendProcessor(ch);
    s->Prepare(var_h, time_out);
T
typhoonzero 已提交
56
    s->response_call_back_ = nullptr;
57

58 59
    auto call = s->stub_g_.PrepareUnaryCall(
        s->context_.get(), "/sendrecv.SendRecvService/SendVariable", req, &cq_);
60
    call->StartCall();
Y
Yi Wang 已提交
61
    call->Finish(&s->reply_, &s->status_, reinterpret_cast<void*>(s));
62
  });
G
gongweibao 已提交
63 64 65 66 67 68
  req_count_++;

  return true;
}

void ProcGetResponse(const VarHandle& var_h,
69
                     const ::grpc::ByteBuffer& ret_msg) {
T
typhoonzero 已提交
70
  framework::Variable* outvar = nullptr;
Y
Yi Wang 已提交
71
  DeserializeFromByteBuffer(ret_msg, *var_h.ctx, var_h.scope, &outvar);
72 73 74 75 76
}

template <typename T>
void RequestToByteBuffer(const T& proto, ::grpc::ByteBuffer* result) {
  ::grpc::Slice slice(proto.ByteSizeLong());
Q
qiaolongfei 已提交
77
  proto.SerializeWithCachedSizesToArray(const_cast<uint8_t*>(slice.begin()));
78 79
  ::grpc::ByteBuffer tmp(&slice, 1);
  result->Swap(&tmp);
G
gongweibao 已提交
80 81 82 83 84 85 86
}

bool RPCClient::AsyncGetVariable(const std::string& ep,
                                 const platform::DeviceContext& ctx,
                                 const framework::Scope& scope,
                                 const std::string& var_name,
                                 int64_t time_out) {
87 88 89 90
  const platform::DeviceContext* p_ctx = &ctx;
  const std::string ep_val = ep;
  const std::string var_name_val = var_name;
  const framework::Scope* p_scope = &scope;
Y
Yancey1989 已提交
91
  const auto ch = GetChannel(ep_val);
92

T
wip  
typhoonzero 已提交
93 94
  framework::AsyncIO([var_name_val, ep_val, p_scope, p_ctx, time_out, ch,
                      this] {
Q
Qiao Longfei 已提交
95
    // prepare input
96 97
    sendrecv::VariableMessage req;
    req.set_varname(var_name_val);
Q
Qiao Longfei 已提交
98 99
    ::grpc::ByteBuffer buf;
    RequestToByteBuffer<sendrecv::VariableMessage>(req, &buf);
100

Q
Qiao Longfei 已提交
101
    // var handle
102 103 104 105 106 107 108 109 110 111 112
    VarHandle var_h;
    var_h.ep = ep_val;
    var_h.scope = p_scope;
    var_h.name = var_name_val;
    var_h.ctx = p_ctx;

    // stub context
    GetProcessor* s = new GetProcessor(ch);
    s->Prepare(var_h, time_out);
    s->response_call_back_ = ProcGetResponse;

113 114
    auto call = s->stub_g_.PrepareUnaryCall(
        s->context_.get(), "/sendrecv.SendRecvService/GetVariable", buf, &cq_);
115
    call->StartCall();
Y
Yi Wang 已提交
116
    call->Finish(&s->reply_, &s->status_, reinterpret_cast<void*>(s));
117
  });
G
gongweibao 已提交
118 119 120 121 122 123

  req_count_++;

  return true;
}

Q
Qiao Longfei 已提交
124 125 126 127 128 129 130 131 132 133 134
bool RPCClient::AsyncPrefetchVariable(const std::string& ep,
                                      const platform::DeviceContext& ctx,
                                      const framework::Scope& scope,
                                      const std::string& in_var_name,
                                      const std::string& out_var_name,
                                      int64_t time_out) {
  const platform::DeviceContext* p_ctx = &ctx;
  const std::string ep_val = ep;
  const std::string in_var_name_val = in_var_name;
  const std::string out_var_name_val = out_var_name;
  const framework::Scope* p_scope = &scope;
Y
Yancey1989 已提交
135
  const auto ch = GetChannel(ep_val);
Q
Qiao Longfei 已提交
136

T
wip  
typhoonzero 已提交
137 138
  framework::AsyncIO([in_var_name_val, out_var_name_val, ep_val, p_scope, p_ctx,
                      time_out, ch, this] {
Q
Qiao Longfei 已提交
139 140 141
    auto* var = p_scope->FindVar(in_var_name_val);

    ::grpc::ByteBuffer req;
Y
Yancey1989 已提交
142
    SerializeToByteBuffer(in_var_name_val, var, *p_ctx, &req, out_var_name_val);
Q
Qiao Longfei 已提交
143 144 145 146 147 148 149 150 151 152 153 154 155 156

    // var handle
    VarHandle var_h;
    var_h.ep = ep_val;
    var_h.scope = p_scope;
    var_h.name = out_var_name_val;
    var_h.ctx = p_ctx;

    // stub context
    GetProcessor* s = new GetProcessor(ch);
    s->Prepare(var_h, time_out);
    s->response_call_back_ = ProcGetResponse;

    auto call = s->stub_g_.PrepareUnaryCall(
157 158
        s->context_.get(), "/sendrecv.SendRecvService/PrefetchVariable", req,
        &cq_);
Q
Qiao Longfei 已提交
159
    call->StartCall();
160
    call->Finish(&s->reply_, &s->status_, static_cast<void*>(s));
Q
Qiao Longfei 已提交
161 162 163 164 165 166
  });

  req_count_++;
  return true;
}

167
void RPCClient::AsyncSendBatchBarrier(const std::string& ep, int64_t time_out) {
Y
Yancey1989 已提交
168
  const auto ch = GetChannel(ep);
Y
Yancey 已提交
169 170 171 172 173 174 175

  BatchBarrierProcessor* s = new BatchBarrierProcessor(ch);
  s->Prepare(time_out);

  sendrecv::VariableMessage req;
  req.set_varname(BATCH_BARRIER_MESSAGE);
  auto rpc = s->stub_->AsyncSendVariable(s->context_.get(), req, &cq_);
Y
Yi Wang 已提交
176
  rpc->Finish(&s->reply_, &s->status_, reinterpret_cast<void*>(s));
Y
Yancey 已提交
177
  req_count_++;
178
}
Y
Yancey 已提交
179

180
void RPCClient::AsyncSendFetchBarrier(const std::string& ep, int64_t time_out) {
Y
Yancey1989 已提交
181
  const auto ch = GetChannel(ep);
182 183 184 185 186 187
  FetchBarrierProcessor* s = new FetchBarrierProcessor(ch);
  s->Prepare(time_out);

  sendrecv::VariableMessage req;
  req.set_varname(FETCH_BARRIER_MESSAGE);
  auto rpc = s->stub_->AsyncGetVariable(s->context_.get(), req, &cq_);
Y
Yi Wang 已提交
188
  rpc->Finish(&s->reply_, &s->status_, reinterpret_cast<void*>(s));
189
  req_count_++;
Y
Yancey 已提交
190 191
}

T
typhoonzero 已提交
192
bool RPCClient::Wait() {
193 194 195
  if (req_count_ <= 0) {
    return true;
  }
Y
Yancey 已提交
196 197
  const size_t kReqCnt = req_count_;
  bool a[kReqCnt];
198
  std::vector<std::future<void>> waits(req_count_);
X
Xin Pan 已提交
199
  std::mutex mu;
G
gongweibao 已提交
200

201
  for (int i = 0; i < req_count_; i++) {
X
Xin Pan 已提交
202 203 204 205 206
    waits[i] = framework::AsyncIO([i, &a, &mu, this] {
      bool ret = Proceed();
      std::lock_guard<std::mutex> l(mu);
      a[i] = ret;
    });
207 208 209 210 211 212 213 214 215 216 217
  }

  for (int i = 0; i < req_count_; i++) {
    waits[i].wait();
  }

  int last_req_count = req_count_;
  req_count_ = 0;

  for (int i = 0; i < last_req_count; i++) {
    if (!a[i]) {
G
gongweibao 已提交
218 219 220 221
      return false;
    }
  }

222
  return true;
G
gongweibao 已提交
223 224 225 226 227 228 229 230
}

bool RPCClient::Proceed() {
  void* tag = NULL;
  bool ok = false;

  // request counts.
  if (!cq_.Next(&tag, &ok)) {
G
gongweibao 已提交
231
    LOG(ERROR) << "Get meets CompletionQueue error";
G
gongweibao 已提交
232 233 234 235 236 237 238
    return false;
  }

  GPR_ASSERT(ok);
  PADDLE_ENFORCE(tag);

  // TODO(gongwb): add more retries.
239
  BaseProcessor* c = static_cast<BaseProcessor*>(tag);
G
gongweibao 已提交
240
  if (!c->status_.ok()) {
G
gongweibao 已提交
241 242
    LOG(ERROR) << "proc param error:" << c->var_h_.String()
               << " grpc error:" << c->status_.error_message();
G
gongweibao 已提交
243
    delete c;
G
gongweibao 已提交
244
    return false;
G
gongweibao 已提交
245 246 247 248 249 250
  }

  c->Process();
  delete c;
  return true;
}
Y
Yancey1989 已提交
251
std::shared_ptr<grpc::Channel> RPCClient::GetChannel(const std::string& ep) {
Y
Yancey1989 已提交
252
  std::unique_lock<std::mutex> lock(mutex_);
Y
Yancey1989 已提交
253
  auto it = channels_.find(ep);
G
gongweibao 已提交
254 255 256 257
  if (it != channels_.end()) {
    return it->second;
  }

G
gongweibao 已提交
258
  grpc::ChannelArguments args;
259
  args.SetCompressionAlgorithm(GRPC_COMPRESS_NONE);
G
gongweibao 已提交
260 261 262
  args.SetMaxSendMessageSize(std::numeric_limits<int>::max());
  args.SetMaxReceiveMessageSize(std::numeric_limits<int>::max());

T
typhoonzero 已提交
263 264
  auto ch =
      grpc::CreateCustomChannel(ep, grpc::InsecureChannelCredentials(), args);
Y
Yancey1989 已提交
265
  channels_[ep] = ch;
G
gongweibao 已提交
266 267 268 269 270 271
  return ch;
}

}  // namespace detail
}  // namespace operators
}  // namespace paddle