“63c74257de0081ed158c3cb48c4eb55b95c3a78a”上不存在“doc/v2/design/cluster_train/data_dispatch.md”
grpc_server.cc 11.3 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. */

15 16
#include <limits>
#include <string>
G
gongweibao 已提交
17

18
#include "paddle/fluid/operators/distributed/grpc_server.h"
G
gongweibao 已提交
19

20
using ::grpc::ServerAsyncResponseWriter;
X
Xin Pan 已提交
21

G
gongweibao 已提交
22 23
namespace paddle {
namespace operators {
24
namespace distributed {
G
gongweibao 已提交
25 26 27 28 29 30
enum CallStatus { PROCESS = 0, FINISH };

// reference:
// https://stackoverflow.com/questions/41732884/grpc-multiple-services-in-cpp-async-server
class RequestBase {
 public:
31
  explicit RequestBase(GrpcService::AsyncService* service,
32 33
                       ::grpc::ServerCompletionQueue* cq,
                       RequestHandler* request_handler, int req_id)
Q
qiaolongfei 已提交
34 35 36
      : service_(service),
        cq_(cq),
        status_(PROCESS),
37 38
        request_handler_(request_handler),
        req_id_(req_id) {
G
gongweibao 已提交
39 40
    PADDLE_ENFORCE(cq_);
  }
G
gongweibao 已提交
41
  virtual ~RequestBase() {}
42
  virtual void Process() = 0;
G
gongweibao 已提交
43

X
Xin Pan 已提交
44 45 46 47 48 49 50 51 52 53 54 55
  CallStatus Status() const {
    std::lock_guard<std::mutex> l(status_mu_);
    return status_;
  }

  template <typename T>
  void Finish(const T& reply, ServerAsyncResponseWriter<T>* responder) {
    std::lock_guard<std::mutex> l(status_mu_);
    status_ = FINISH;
    responder->Finish(reply, ::grpc::Status::OK,
                      reinterpret_cast<void*>(static_cast<intptr_t>(req_id_)));
  }
56
  virtual std::string GetReqName() = 0;
G
gongweibao 已提交
57 58

 protected:
X
Xin Pan 已提交
59
  mutable std::mutex status_mu_;
60 61 62
  ::grpc::ServerContext ctx_;
  GrpcService::AsyncService* service_;
  ::grpc::ServerCompletionQueue* cq_;
G
gongweibao 已提交
63
  CallStatus status_;
64 65
  RequestHandler* request_handler_;
  int req_id_;
G
gongweibao 已提交
66 67 68 69
};

class RequestSend final : public RequestBase {
 public:
70
  explicit RequestSend(GrpcService::AsyncService* service,
71 72 73 74 75 76
                       ::grpc::ServerCompletionQueue* cq,
                       RequestHandler* request_handler, int req_id)
      : RequestBase(service, cq, request_handler, req_id), responder_(&ctx_) {
    request_.reset(new VariableResponse(request_handler->scope(),
                                        request_handler->dev_ctx(),
                                        !request_handler->sync_mode()));
77
    int method_id = static_cast<int>(distributed::GrpcMethod::kSendVariable);
X
Xin Pan 已提交
78 79
    service_->RequestAsyncUnary(
        method_id, &ctx_, request_.get(), &responder_, cq_, cq_,
X
Xin Pan 已提交
80
        reinterpret_cast<void*>(static_cast<intptr_t>(req_id)));
G
gongweibao 已提交
81 82
  }
  virtual ~RequestSend() {}
83 84 85 86 87
  std::string GetReqName() override { return request_->Varname(); }

  void Process() override {
    std::string varname = GetReqName();
    VLOG(3) << "RequestSend var_name:" << varname;
G
gongweibao 已提交
88

89 90 91 92 93
    auto scope = request_->GetMutableLocalScope();
    auto invar = request_->GetVar();
    framework::Variable* outvar = nullptr;

    request_handler_->Handle(varname, scope, invar, &outvar);
X
Xin Pan 已提交
94
    Finish(reply_, &responder_);
G
gongweibao 已提交
95 96 97
  }

 protected:
X
Xin Pan 已提交
98
  sendrecv::VoidMessage reply_;
99
  std::shared_ptr<VariableResponse> request_;
G
gongweibao 已提交
100 101 102 103 104
  ServerAsyncResponseWriter<sendrecv::VoidMessage> responder_;
};

class RequestGet final : public RequestBase {
 public:
105
  explicit RequestGet(GrpcService::AsyncService* service,
106 107 108
                      ::grpc::ServerCompletionQueue* cq,
                      RequestHandler* request_handler, int req_id)
      : RequestBase(service, cq, request_handler, req_id), responder_(&ctx_) {
109
    auto method_id = static_cast<int>(distributed::GrpcMethod::kGetVariable);
X
Xin Pan 已提交
110 111
    service_->RequestAsyncUnary(
        method_id, &ctx_, &request_, &responder_, cq_, cq_,
112
        reinterpret_cast<void*>(static_cast<intptr_t>(req_id)));
G
gongweibao 已提交
113 114 115 116
  }

  virtual ~RequestGet() {}

117
  std::string GetReqName() override { return request_.varname(); }
G
gongweibao 已提交
118

119
  void Process() override {
G
gongweibao 已提交
120
    // proc request.
121 122 123 124 125 126
    std::string varname = request_.varname();
    VLOG(3) << "RequestGet " << varname;

    auto scope = request_handler_->scope();
    auto invar = scope->FindVar(varname);
    framework::Variable* outvar = nullptr;
127

128 129 130 131 132
    request_handler_->Handle(varname, scope, invar, &outvar);

    if (outvar) {
      SerializeToByteBuffer(varname, outvar, *request_handler_->dev_ctx(),
                            &reply_);
133
    }
X
Xin Pan 已提交
134
    Finish(reply_, &responder_);
G
gongweibao 已提交
135 136 137 138
  }

 protected:
  sendrecv::VariableMessage request_;
X
Xin Pan 已提交
139
  ::grpc::ByteBuffer reply_;
140
  ServerAsyncResponseWriter<::grpc::ByteBuffer> responder_;
G
gongweibao 已提交
141 142
};

143 144 145
class RequestPrefetch final : public RequestBase {
 public:
  explicit RequestPrefetch(GrpcService::AsyncService* service,
146 147 148
                           ::grpc::ServerCompletionQueue* cq,
                           RequestHandler* request_handler, int req_id)
      : RequestBase(service, cq, request_handler, req_id),
149
        responder_(&ctx_),
150 151 152
        local_scope_(nullptr) {
    request_.reset(new VariableResponse(request_handler->scope(),
                                        request_handler->dev_ctx(), true));
153 154
    int method_id =
        static_cast<int>(distributed::GrpcMethod::kPrefetchVariable);
X
Xin Pan 已提交
155 156
    service_->RequestAsyncUnary(
        method_id, &ctx_, request_.get(), &responder_, cq_, cq_,
157
        reinterpret_cast<void*>(static_cast<intptr_t>(req_id)));
158 159 160 161
  }

  virtual ~RequestPrefetch() {}

162
  std::string GetReqName() override { return request_->Varname(); }
163

164
  void Process() override {
165
    // prefetch process...
166 167
    std::string in_var_name = request_->Varname();
    std::string out_var_name = request_->OutVarname();
168 169
    VLOG(3) << "RequestPrefetch, in_var_name: " << in_var_name
            << " out_var_name: " << out_var_name;
170 171

    auto scope = request_->GetMutableLocalScope();
172
    auto invar = scope->FindVar(in_var_name);
173
    // out var must be created in local scope!
Q
qiaolongfei 已提交
174
    framework::Variable* outvar = scope->Var(out_var_name);
175

Q
qiaolongfei 已提交
176
    request_handler_->Handle(in_var_name, scope, invar, &outvar, out_var_name);
Y
Yancey1989 已提交
177

178
    SerializeToByteBuffer(out_var_name, outvar, *request_handler_->dev_ctx(),
179
                          &reply_);
X
Xin Pan 已提交
180
    Finish(reply_, &responder_);
181 182 183
  }

 protected:
Y
Yancey1989 已提交
184
  std::shared_ptr<VariableResponse> request_;
X
Xin Pan 已提交
185
  ::grpc::ByteBuffer reply_;
186
  ServerAsyncResponseWriter<::grpc::ByteBuffer> responder_;
187
  framework::Scope* local_scope_;
188 189
};

T
done  
typhoonzero 已提交
190
void AsyncGRPCServer::WaitServerReady() {
191
  VLOG(3) << "AsyncGRPCServer is wait server ready";
T
update  
typhoonzero 已提交
192
  std::unique_lock<std::mutex> lock(this->mutex_ready_);
T
done  
typhoonzero 已提交
193
  condition_ready_.wait(lock, [=] { return this->ready_ == 1; });
194
  VLOG(3) << "AsyncGRPCServer WaitSeverReady";
T
update  
typhoonzero 已提交
195 196
}

197
void AsyncGRPCServer::StartServer() {
198
  ::grpc::ServerBuilder builder;
199
  builder.AddListeningPort(bind_address_, ::grpc::InsecureServerCredentials(),
T
typhoonzero 已提交
200
                           &selected_port_);
201

G
gongweibao 已提交
202 203
  builder.SetMaxSendMessageSize(std::numeric_limits<int>::max());
  builder.SetMaxReceiveMessageSize(std::numeric_limits<int>::max());
G
gongweibao 已提交
204 205
  builder.RegisterService(&service_);

206 207 208
  for (auto t : rpc_call_map_) {
    rpc_cq_[t.first].reset(builder.AddCompletionQueue().release());
  }
Y
Yancey 已提交
209

G
gongweibao 已提交
210
  server_ = builder.BuildAndStart();
211
  LOG(INFO) << "Server listening on " << bind_address_
T
typhoonzero 已提交
212
            << " selected port: " << selected_port_;
G
gongweibao 已提交
213

214 215 216
  std::function<void(const std::string&, int)> f =
      std::bind(&AsyncGRPCServer::TryToRegisterNewOne, this,
                std::placeholders::_1, std::placeholders::_2);
X
Xin Pan 已提交
217

218 219 220 221 222
  for (auto& t : rpc_call_map_) {
    auto& rpc_name = t.first;
    auto& cq = rpc_cq_[rpc_name];
    auto threadnum = rpc_thread_num_[rpc_name];
    auto& reqs = rpc_reqs_[rpc_name];
X
Xin Pan 已提交
223

224 225 226 227 228 229 230 231 232 233 234
    reqs.reserve(kRequestBufSize);

    for (int i = 0; i < kRequestBufSize; i++) {
      TryToRegisterNewOne(rpc_name, i);
    }

    for (int i = 0; i < threadnum; i++) {
      rpc_threads_[rpc_name].emplace_back(new std::thread(std::bind(
          &AsyncGRPCServer::HandleRequest, this, cq.get(), rpc_name, f)));
      VLOG(3) << t.first << " creates threads!";
    }
X
Xin Pan 已提交
235
  }
236

T
wip  
typhoonzero 已提交
237 238 239 240 241
  {
    std::lock_guard<std::mutex> lock(this->mutex_ready_);
    ready_ = 1;
  }
  condition_ready_.notify_all();
242

G
gongweibao 已提交
243 244
  // wait server
  server_->Wait();
245 246 247 248 249 250 251

  for (auto& t : rpc_threads_) {
    auto& threads = t.second;
    for (size_t i = 0; i < threads.size(); ++i) {
      threads[i]->join();
      VLOG(3) << t.first << " threads ends!";
    }
X
Xin Pan 已提交
252
  }
G
gongweibao 已提交
253 254 255
}

void AsyncGRPCServer::ShutdownQueue() {
256 257 258 259
  for (auto& t : rpc_cq_) {
    t.second->Shutdown();
    VLOG(3) << t.first << " shutdown!";
  }
G
gongweibao 已提交
260 261
}

262 263
void AsyncGRPCServer::ShutDownImpl() {
  std::unique_lock<std::mutex> lock(cq_mutex_);
T
typhoonzero 已提交
264
  is_shut_down_ = true;
G
gongweibao 已提交
265
  ShutdownQueue();
266 267

  VLOG(3) << "server_ shutdown!";
T
typhoonzero 已提交
268
  server_->Shutdown();
G
gongweibao 已提交
269 270
}

271 272
void AsyncGRPCServer::TryToRegisterNewOne(const std::string& rpc_name,
                                          int req_id) {
G
gongweibao 已提交
273 274
  std::unique_lock<std::mutex> lock(cq_mutex_);
  if (is_shut_down_) {
275
    VLOG(3) << "shutdown, do not TryToRegisterNewSendOne";
G
gongweibao 已提交
276 277 278
    return;
  }

279 280 281 282 283 284 285 286 287 288 289 290 291 292 293
  VLOG(4) << "register send rpc_name:" << rpc_name
          << ", handler:" << rpc_call_map_[kRequestSend];

  auto& reqs = rpc_reqs_[rpc_name];
  auto& handler = rpc_call_map_[rpc_name];
  auto& cq = rpc_cq_[rpc_name];

  RequestBase* b = nullptr;
  if (rpc_name == kRequestSend) {
    b = new RequestSend(&service_, cq.get(), handler, req_id);
  } else if (rpc_name == kRequestGet) {
    b = new RequestGet(&service_, cq.get(), handler, req_id);
  } else if (rpc_name == kRequestPrefetch) {
    b = new RequestPrefetch(&service_, cq.get(), handler, req_id);
  } else {
Q
qiaolongfei 已提交
294
    PADDLE_ENFORCE(false, "not supported rpc");
G
gongweibao 已提交
295 296
  }

297
  reqs[req_id] = b;
298

299
  VLOG(4) << "Create RequestSend status:" << b->Status();
300 301
}

X
Xin Pan 已提交
302
void AsyncGRPCServer::HandleRequest(
303 304
    ::grpc::ServerCompletionQueue* cq, const std::string& rpc_name,
    std::function<void(const std::string&, int)> TryToRegisterNewOne) {
G
gongweibao 已提交
305 306
  void* tag = NULL;
  bool ok = false;
307

G
gongweibao 已提交
308
  while (true) {
309
    VLOG(3) << "HandleRequest " << rpc_name << " wait next";
G
gongweibao 已提交
310
    if (!cq->Next(&tag, &ok)) {
311
      LOG(INFO) << "CompletionQueue " << rpc_name << " shutdown!";
G
gongweibao 已提交
312 313
      break;
    }
Q
qiaolongfei 已提交
314

315 316 317
    int req_id = static_cast<int>(reinterpret_cast<intptr_t>(tag));
    VLOG(3) << "HandleRequest " << rpc_name << ", req_id:" << req_id
            << " get next";
G
gongweibao 已提交
318

319
    auto& reqs = rpc_reqs_[rpc_name];
X
Xin Pan 已提交
320 321
    RequestBase* base = nullptr;
    {
322 323 324
      PADDLE_ENFORCE(req_id >= 0 && req_id < kRequestBufSize);
      std::unique_lock<std::mutex> lock(cq_mutex_);
      base = reqs[req_id];
X
Xin Pan 已提交
325
    }
326

G
gongweibao 已提交
327 328 329 330
    // reference:
    // https://github.com/tensorflow/tensorflow/issues/5596
    // https://groups.google.com/forum/#!topic/grpc-io/xftlRy-IQwM
    // https://groups.google.com/forum/#!topic/grpc-io/ywATt88Ef_I
G
gongweibao 已提交
331
    if (!ok) {
332 333
      LOG(WARNING) << "completion queue:" << rpc_name
                   << " recv no regular event:argument name["
Q
qiaolongfei 已提交
334
                   << base->GetReqName() << "]";
335
      TryToRegisterNewOne(rpc_name, req_id);
G
gongweibao 已提交
336 337 338 339
      delete base;
      continue;
    }

340 341 342
    VLOG(3) << "queue id:" << rpc_name << ", req_id:" << req_id
            << ", status:" << base->Status();

G
gongweibao 已提交
343 344 345 346 347 348
    switch (base->Status()) {
      case PROCESS: {
        base->Process();
        break;
      }
      case FINISH: {
349
        TryToRegisterNewOne(rpc_name, req_id);
G
gongweibao 已提交
350 351 352 353 354 355 356 357
        delete base;
        break;
      }
      default: { assert(false); }
    }
  }
}

358
}  // namespace distributed
G
gongweibao 已提交
359 360
}  // namespace operators
}  // namespace paddle