grpc_client.cc 5.2 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 15

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 "grpc_client.h"
Y
Yi Wang 已提交
16
#include "paddle/fluid/framework/threadpool.h"
G
gongweibao 已提交
17 18 19 20 21 22 23 24 25
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) {
26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51
  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;
  const auto ch = GetChannel(ep_val);

  framework::Async([var_name_val, p_ctx, ep_val, p_scope, time_out, ch, this] {
    auto* var = p_scope->FindVar(var_name_val);
    sendrecv::VariableMessage req;
    SerializeToMessage(var_name_val, var, *p_ctx, &req);

    // 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);
    s->response_call_back_ = NULL;

    auto rpc = s->stub_->AsyncSendVariable(s->context_.get(), req, &cq_);
    rpc->Finish(&s->reply_, &s->status_, (void*)s);
  });
G
gongweibao 已提交
52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68

  req_count_++;

  return true;
}

void ProcGetResponse(const VarHandle& var_h,
                     const sendrecv::VariableMessage& ret_msg) {
  auto* outvar = var_h.scope->FindVar(var_h.name);
  DeserializeFromMessage(ret_msg, *var_h.ctx, outvar);
}

bool RPCClient::AsyncGetVariable(const std::string& ep,
                                 const platform::DeviceContext& ctx,
                                 const framework::Scope& scope,
                                 const std::string& var_name,
                                 int64_t time_out) {
69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93
  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;
  const auto ch = GetChannel(ep_val);

  framework::Async([var_name_val, ep_val, p_scope, p_ctx, time_out, ch, this] {
    sendrecv::VariableMessage req;
    req.set_varname(var_name_val);

    // 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
    GetProcessor* s = new GetProcessor(ch);
    s->Prepare(var_h, time_out);
    s->response_call_back_ = ProcGetResponse;

    auto rpc = s->stub_->AsyncGetVariable(s->context_.get(), req, &cq_);
    rpc->Finish(&s->reply_, &s->status_, (void*)s);
  });
G
gongweibao 已提交
94 95 96 97 98 99

  req_count_++;

  return true;
}

Y
Yancey 已提交
100 101 102 103 104 105 106 107 108 109 110 111 112 113 114
bool RPCClient::AsyncSendBatchBarrier(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(BATCH_BARRIER_MESSAGE);
  auto rpc = s->stub_->AsyncSendVariable(s->context_.get(), req, &cq_);
  rpc->Finish(&s->reply_, &s->status_, (void*)s);
  req_count_++;

  return true;
}

T
typhoonzero 已提交
115
bool RPCClient::Wait() {
116 117 118
  if (req_count_ <= 0) {
    return true;
  }
Y
Yancey 已提交
119 120
  const size_t kReqCnt = req_count_;
  bool a[kReqCnt];
121
  std::vector<std::future<void>> waits(req_count_);
G
gongweibao 已提交
122

123 124 125 126 127 128 129 130 131 132 133 134 135
  for (int i = 0; i < req_count_; i++) {
    waits[i] = framework::Async([i, &a, this] { a[i] = Proceed(); });
  }

  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 已提交
136 137 138 139
      return false;
    }
  }

140
  return true;
G
gongweibao 已提交
141 142 143 144 145 146 147 148
}

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

  // request counts.
  if (!cq_.Next(&tag, &ok)) {
G
gongweibao 已提交
149
    LOG(ERROR) << "Get meets CompletionQueue error";
G
gongweibao 已提交
150 151 152 153 154 155 156 157 158
    return false;
  }

  GPR_ASSERT(ok);
  PADDLE_ENFORCE(tag);

  // TODO(gongwb): add more retries.
  ClientBase* c = static_cast<ClientBase*>(tag);
  if (!c->status_.ok()) {
G
gongweibao 已提交
159 160
    LOG(ERROR) << "proc param error:" << c->var_h_.String()
               << " grpc error:" << c->status_.error_message();
G
gongweibao 已提交
161
    delete c;
G
gongweibao 已提交
162
    return false;
G
gongweibao 已提交
163 164 165 166 167 168 169 170 171 172 173 174 175
  }

  c->Process();
  delete c;
  return true;
}

std::shared_ptr<grpc::Channel> RPCClient::GetChannel(const std::string& ep) {
  auto it = channels_.find(ep);
  if (it != channels_.end()) {
    return it->second;
  }

G
gongweibao 已提交
176 177 178 179
  grpc::ChannelArguments args;
  args.SetMaxSendMessageSize(std::numeric_limits<int>::max());
  args.SetMaxReceiveMessageSize(std::numeric_limits<int>::max());

T
typhoonzero 已提交
180 181
  auto ch =
      grpc::CreateCustomChannel(ep, grpc::InsecureChannelCredentials(), args);
G
gongweibao 已提交
182 183 184 185 186 187 188 189

  channels_[ep] = ch;
  return ch;
}

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