grpc_client.cc 8.5 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"
G
gongweibao 已提交
22
#include "paddle/fluid/operators/detail/request_handler.h"
X
Xin Pan 已提交
23
#include "paddle/fluid/platform/profiler.h"
24

G
gongweibao 已提交
25 26 27 28
namespace paddle {
namespace operators {
namespace detail {

G
gongweibao 已提交
29
void GRPCClient::InitImpl() { InitEventLoop(); }
Y
Yancey1989 已提交
30

G
gongweibao 已提交
31
void GRPCClient::InitEventLoop() {
W
Wu Yi 已提交
32 33
  // start the client process thread
  // TODO(wuyi): can make this in a threadpool
G
gongweibao 已提交
34
  client_thread_.reset(new std::thread(std::bind(&GRPCClient::Proceed, this)));
W
Wu Yi 已提交
35 36
}

W
Wu Yi 已提交
37 38 39 40 41 42
void GRPCClient::SendComplete() {
  for (auto& it : channels_) {
    this->AsyncSendComplete(it.first);
  }
}

G
gongweibao 已提交
43
GRPCClient::~GRPCClient() {
W
Wu Yi 已提交
44 45 46 47 48 49 50 51 52
  Wait();
  cq_.Shutdown();
  {
    std::lock_guard<std::mutex> guard(chan_mutex_);
    for (auto& it : channels_) {
      it.second.reset();
    }
  }
  client_thread_->join();
Y
Yancey1989 已提交
53 54
}

G
gongweibao 已提交
55 56 57 58
bool GRPCClient::AsyncSendVar(const std::string& ep,
                              const platform::DeviceContext& ctx,
                              const framework::Scope& scope,
                              const std::string& var_name, int64_t time_out) {
59 60 61 62
  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 已提交
63
  const auto ch = GetChannel(ep_val);
64

T
wip  
typhoonzero 已提交
65 66
  framework::AsyncIO([var_name_val, p_ctx, ep_val, p_scope, time_out, ch,
                      this] {
67
    auto* var = p_scope->FindVar(var_name_val);
68 69 70

    ::grpc::ByteBuffer req;
    SerializeToByteBuffer(var_name_val, var, *p_ctx, &req);
71 72 73 74 75 76 77 78 79 80 81

    // 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 已提交
82
    s->response_call_back_ = nullptr;
83

84 85
    auto call = s->stub_g_.PrepareUnaryCall(
        s->context_.get(), "/sendrecv.SendRecvService/SendVariable", req, &cq_);
86
    call->StartCall();
Y
Yi Wang 已提交
87
    call->Finish(&s->reply_, &s->status_, reinterpret_cast<void*>(s));
88
  });
G
gongweibao 已提交
89 90 91 92 93 94
  req_count_++;

  return true;
}

void ProcGetResponse(const VarHandle& var_h,
95
                     const ::grpc::ByteBuffer& ret_msg) {
T
typhoonzero 已提交
96
  framework::Variable* outvar = nullptr;
Y
Yi Wang 已提交
97
  DeserializeFromByteBuffer(ret_msg, *var_h.ctx, var_h.scope, &outvar);
98 99 100 101 102
}

template <typename T>
void RequestToByteBuffer(const T& proto, ::grpc::ByteBuffer* result) {
  ::grpc::Slice slice(proto.ByteSizeLong());
Q
qiaolongfei 已提交
103
  proto.SerializeWithCachedSizesToArray(const_cast<uint8_t*>(slice.begin()));
104 105
  ::grpc::ByteBuffer tmp(&slice, 1);
  result->Swap(&tmp);
G
gongweibao 已提交
106 107
}

G
gongweibao 已提交
108 109 110 111
bool GRPCClient::AsyncGetVar(const std::string& ep,
                             const platform::DeviceContext& ctx,
                             const framework::Scope& scope,
                             const std::string& var_name, int64_t time_out) {
112 113 114 115
  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 已提交
116
  const auto ch = GetChannel(ep_val);
117

T
wip  
typhoonzero 已提交
118 119
  framework::AsyncIO([var_name_val, ep_val, p_scope, p_ctx, time_out, ch,
                      this] {
Q
Qiao Longfei 已提交
120
    // prepare input
121 122
    sendrecv::VariableMessage req;
    req.set_varname(var_name_val);
Q
Qiao Longfei 已提交
123 124
    ::grpc::ByteBuffer buf;
    RequestToByteBuffer<sendrecv::VariableMessage>(req, &buf);
125

Q
Qiao Longfei 已提交
126
    // var handle
127 128 129 130 131 132 133 134 135 136 137
    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;

138 139
    auto call = s->stub_g_.PrepareUnaryCall(
        s->context_.get(), "/sendrecv.SendRecvService/GetVariable", buf, &cq_);
140
    call->StartCall();
Y
Yi Wang 已提交
141
    call->Finish(&s->reply_, &s->status_, reinterpret_cast<void*>(s));
142
  });
G
gongweibao 已提交
143 144 145 146 147 148

  req_count_++;

  return true;
}

G
gongweibao 已提交
149 150 151 152 153 154
bool GRPCClient::AsyncPrefetchVar(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) {
Q
Qiao Longfei 已提交
155 156 157 158 159
  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 已提交
160
  const auto ch = GetChannel(ep_val);
Q
Qiao Longfei 已提交
161

T
wip  
typhoonzero 已提交
162 163
  framework::AsyncIO([in_var_name_val, out_var_name_val, ep_val, p_scope, p_ctx,
                      time_out, ch, this] {
Q
Qiao Longfei 已提交
164 165 166
    auto* var = p_scope->FindVar(in_var_name_val);

    ::grpc::ByteBuffer req;
Y
Yancey1989 已提交
167
    SerializeToByteBuffer(in_var_name_val, var, *p_ctx, &req, out_var_name_val);
Q
Qiao Longfei 已提交
168 169 170 171 172 173 174 175 176 177 178 179 180 181

    // 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(
182 183
        s->context_.get(), "/sendrecv.SendRecvService/PrefetchVariable", req,
        &cq_);
Q
Qiao Longfei 已提交
184
    call->StartCall();
185
    call->Finish(&s->reply_, &s->status_, static_cast<void*>(s));
Q
Qiao Longfei 已提交
186 187 188 189 190 191
  });

  req_count_++;
  return true;
}

G
gongweibao 已提交
192 193
void GRPCClient::AsyncSendBatchBarrier(const std::string& ep,
                                       int64_t time_out) {
Y
Yancey1989 已提交
194
  const auto ch = GetChannel(ep);
Y
Yancey 已提交
195 196 197 198 199 200 201

  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 已提交
202
  rpc->Finish(&s->reply_, &s->status_, reinterpret_cast<void*>(s));
Y
Yancey 已提交
203
  req_count_++;
204
}
Y
Yancey 已提交
205

G
gongweibao 已提交
206 207
void GRPCClient::AsyncSendFetchBarrier(const std::string& ep,
                                       int64_t time_out) {
Y
Yancey1989 已提交
208
  const auto ch = GetChannel(ep);
209 210 211 212 213 214
  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 已提交
215
  rpc->Finish(&s->reply_, &s->status_, reinterpret_cast<void*>(s));
216
  req_count_++;
Y
Yancey 已提交
217 218
}

W
Wu Yi 已提交
219 220 221 222 223 224 225 226 227 228 229 230 231
void GRPCClient::AsyncSendComplete(const std::string& ep, int64_t time_out) {
  const auto ch = GetChannel(ep);

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

  sendrecv::VariableMessage req;
  req.set_varname(COMPLETE_MESSAGE);
  auto rpc = s->stub_->AsyncSendVariable(s->context_.get(), req, &cq_);
  rpc->Finish(&s->reply_, &s->status_, reinterpret_cast<void*>(s));
  req_count_++;
}

G
gongweibao 已提交
232
void GRPCClient::Wait() {
W
Wu Yi 已提交
233 234
  std::unique_lock<std::mutex> lk(sync_mutex_);
  sync_cond_.wait(lk, [this] { return req_count_ == 0; });
G
gongweibao 已提交
235 236
}

G
gongweibao 已提交
237
void GRPCClient::Proceed() {
W
Wu Yi 已提交
238
  void* tag = nullptr;
G
gongweibao 已提交
239 240
  bool ok = false;

W
Wu Yi 已提交
241 242 243 244 245 246 247 248 249 250
  while (cq_.Next(&tag, &ok)) {
    BaseProcessor* c = static_cast<BaseProcessor*>(tag);
    GPR_ASSERT(ok);
    PADDLE_ENFORCE(c);
    if (c->status_.ok()) {
      c->Process();
    } else {
      LOG(ERROR) << "var: " << c->var_h_.String()
                 << " grpc error:" << c->status_.error_message();
    }
G
gongweibao 已提交
251
    delete c;
W
Wu Yi 已提交
252 253 254 255 256
    {
      std::lock_guard<std::mutex> lk(sync_mutex_);
      req_count_--;
    }
    sync_cond_.notify_all();
G
gongweibao 已提交
257 258
  }
}
W
Wu Yi 已提交
259

G
gongweibao 已提交
260
std::shared_ptr<grpc::Channel> GRPCClient::GetChannel(const std::string& ep) {
Y
Yancey1989 已提交
261
  // TODO(Yancey1989): make grpc client completely thread-safe
W
Wu Yi 已提交
262
  std::lock_guard<std::mutex> guard(chan_mutex_);
Y
Yancey1989 已提交
263
  auto it = channels_.find(ep);
G
gongweibao 已提交
264 265 266 267
  if (it != channels_.end()) {
    return it->second;
  }

G
gongweibao 已提交
268
  grpc::ChannelArguments args;
269
  args.SetCompressionAlgorithm(GRPC_COMPRESS_NONE);
G
gongweibao 已提交
270 271 272
  args.SetMaxSendMessageSize(std::numeric_limits<int>::max());
  args.SetMaxReceiveMessageSize(std::numeric_limits<int>::max());

T
typhoonzero 已提交
273 274
  auto ch =
      grpc::CreateCustomChannel(ep, grpc::InsecureChannelCredentials(), args);
Y
Yancey1989 已提交
275
  channels_[ep] = ch;
G
gongweibao 已提交
276 277 278 279 280 281
  return ch;
}

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