grpc_server.cc 24.1 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
#include <unistd.h>
16
#include <limits>
Q
Qiao Longfei 已提交
17
#include <memory>
18
#include <string>
G
gongweibao 已提交
19

W
Wu Yi 已提交
20 21
#include "paddle/fluid/operators/distributed/grpc/grpc_serde.h"
#include "paddle/fluid/operators/distributed/grpc/grpc_server.h"
G
gongweibao 已提交
22

W
wanghuancoder 已提交
23 24 25 26 27 28 29 30 31 32 33 34 35 36
namespace grpc {
class ChannelArguments;
}  // namespace grpc
namespace paddle {
namespace framework {
class Variable;
}  // namespace framework
namespace operators {
namespace distributed {
class GRPCVariableResponse;
}  // namespace distributed
}  // namespace operators
}  // namespace paddle

37
using ::grpc::ServerAsyncResponseWriter;
X
Xin Pan 已提交
38

39
DECLARE_bool(rpc_disable_reuse_port);
40
DECLARE_int32(rpc_retry_bind_port);
41

G
gongweibao 已提交
42 43
namespace paddle {
namespace operators {
44
namespace distributed {
45

G
gongweibao 已提交
46 47 48 49 50 51
enum CallStatus { PROCESS = 0, FINISH };

// reference:
// https://stackoverflow.com/questions/41732884/grpc-multiple-services-in-cpp-async-server
class RequestBase {
 public:
52
  explicit RequestBase(GrpcService::AsyncService* service,
53 54
                       ::grpc::ServerCompletionQueue* cq,
                       RequestHandler* request_handler, int req_id)
Q
qiaolongfei 已提交
55 56 57
      : service_(service),
        cq_(cq),
        status_(PROCESS),
58 59
        request_handler_(request_handler),
        req_id_(req_id) {
G
gongweibao 已提交
60 61
    PADDLE_ENFORCE(cq_);
  }
G
gongweibao 已提交
62
  virtual ~RequestBase() {}
63
  virtual void Process() = 0;
G
gongweibao 已提交
64

G
gongweibao 已提交
65 66 67 68 69 70 71 72 73 74 75 76 77
  std::string Status2String(const std::string& method) {
    std::string status = "Process";
    if (status_ == FINISH) {
      status = "Finish";
    }

    std::ostringstream s;
    s << method << " name:[" << GetReqName() << "]"
      << ", ep:[" << ctx_.peer() << "]"
      << " " << status << " using req_id:" << req_id_;
    return s.str();
  }

X
Xin Pan 已提交
78 79 80 81 82 83 84 85 86 87 88 89
  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_)));
  }
90
  virtual std::string GetReqName() = 0;
G
gongweibao 已提交
91 92

 protected:
X
Xin Pan 已提交
93
  mutable std::mutex status_mu_;
94 95 96
  ::grpc::ServerContext ctx_;
  GrpcService::AsyncService* service_;
  ::grpc::ServerCompletionQueue* cq_;
G
gongweibao 已提交
97
  CallStatus status_;
98 99
  RequestHandler* request_handler_;
  int req_id_;
G
gongweibao 已提交
100 101 102 103
};

class RequestSend final : public RequestBase {
 public:
104
  explicit RequestSend(GrpcService::AsyncService* service,
105 106 107
                       ::grpc::ServerCompletionQueue* cq,
                       RequestHandler* request_handler, int req_id)
      : RequestBase(service, cq, request_handler, req_id), responder_(&ctx_) {
1
123malin 已提交
108 109 110
    request_.reset(new GRPCVariableResponse(
        request_handler->scope(), request_handler->dev_ctx(),
        request_handler->distributed_mode()));
111
    int method_id = static_cast<int>(distributed::GrpcMethod::kSendVariable);
X
Xin Pan 已提交
112 113
    service_->RequestAsyncUnary(
        method_id, &ctx_, request_.get(), &responder_, cq_, cq_,
X
Xin Pan 已提交
114
        reinterpret_cast<void*>(static_cast<intptr_t>(req_id)));
G
gongweibao 已提交
115 116
  }
  virtual ~RequestSend() {}
117 118 119 120
  std::string GetReqName() override { return request_->Varname(); }

  void Process() override {
    std::string varname = GetReqName();
G
gongweibao 已提交
121

122 123
    auto scope = request_->GetMutableLocalScope();
    auto invar = request_->GetVar();
W
Wu Yi 已提交
124
    int trainer_id = request_->GetTrainerId();
125 126 127

    VLOG(4) << "RequestSend var_name:" << varname << " trainer: " << trainer_id;

128
    framework::Variable* outvar = nullptr;
W
Wu Yi 已提交
129
    request_handler_->Handle(varname, scope, invar, &outvar, trainer_id);
X
Xin Pan 已提交
130
    Finish(reply_, &responder_);
G
gongweibao 已提交
131 132 133
  }

 protected:
X
Xin Pan 已提交
134
  sendrecv::VoidMessage reply_;
135
  std::shared_ptr<GRPCVariableResponse> request_;
G
gongweibao 已提交
136 137 138 139 140
  ServerAsyncResponseWriter<sendrecv::VoidMessage> responder_;
};

class RequestGet final : public RequestBase {
 public:
141
  explicit RequestGet(GrpcService::AsyncService* service,
142 143 144
                      ::grpc::ServerCompletionQueue* cq,
                      RequestHandler* request_handler, int req_id)
      : RequestBase(service, cq, request_handler, req_id), responder_(&ctx_) {
145
    auto method_id = static_cast<int>(distributed::GrpcMethod::kGetVariable);
X
Xin Pan 已提交
146 147
    service_->RequestAsyncUnary(
        method_id, &ctx_, &request_, &responder_, cq_, cq_,
148
        reinterpret_cast<void*>(static_cast<intptr_t>(req_id)));
G
gongweibao 已提交
149 150 151 152
  }

  virtual ~RequestGet() {}

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

155
  void Process() override {
G
gongweibao 已提交
156
    // proc request.
157
    std::string varname = request_.varname();
158
    std::string out_varname = request_.out_varname();
Q
Qiao Longfei 已提交
159
    std::string table_name = request_.table_name();
W
Wu Yi 已提交
160
    int trainer_id = request_.trainer_id();
161 162

    VLOG(4) << "RequestGet " << out_varname << " from " << varname;
163 164

    auto scope = request_handler_->scope();
165
    framework::Variable* invar = nullptr;
166
    framework::Variable* outvar = nullptr;
167

Q
Qiao Longfei 已提交
168
    tmp_scope_ = std::move(scope->NewTmpScope());
Q
Qiao Longfei 已提交
169 170
    request_handler_->Handle(varname, tmp_scope_.get(), invar, &outvar,
                             trainer_id, out_varname, table_name);
171

Q
Qiao Longfei 已提交
172
    VLOG(1) << "before SerializeToByteBuffer";
173
    if (outvar) {
174 175 176
      SerializeToByteBuffer(out_varname, outvar, *request_handler_->dev_ctx(),
                            &reply_);
    }
Q
Qiao Longfei 已提交
177
    VLOG(1) << "after SerializeToByteBuffer";
178 179 180 181 182 183
    Finish(reply_, &responder_);
  }

 protected:
  sendrecv::VariableMessage request_;
  ::grpc::ByteBuffer reply_;
Q
Qiao Longfei 已提交
184
  std::unique_ptr<framework::Scope> tmp_scope_;
185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 204 205 206 207 208 209 210 211 212 213 214 215 216 217 218 219 220 221
  ServerAsyncResponseWriter<::grpc::ByteBuffer> responder_;
};

class RequestGetNoBarrier final : public RequestBase {
 public:
  explicit RequestGetNoBarrier(GrpcService::AsyncService* service,
                               ::grpc::ServerCompletionQueue* cq,
                               RequestHandler* request_handler, int req_id)
      : RequestBase(service, cq, request_handler, req_id), responder_(&ctx_) {
    auto method_id =
        static_cast<int>(distributed::GrpcMethod::kGetVariableNoBarrier);
    service_->RequestAsyncUnary(
        method_id, &ctx_, &request_, &responder_, cq_, cq_,
        reinterpret_cast<void*>(static_cast<intptr_t>(req_id)));
  }

  virtual ~RequestGetNoBarrier() {}

  std::string GetReqName() override { return request_.varname(); }

  void Process() override {
    // proc request.
    std::string varname = request_.varname();
    std::string out_varname = request_.out_varname();
    int trainer_id = request_.trainer_id();

    VLOG(4) << "RequestGetNoBarrier " << out_varname << " from " << varname;

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

    request_handler_->Handle(varname, scope, invar, &outvar, trainer_id,
                             out_varname);

    if (outvar) {
      SerializeToByteBuffer(out_varname, outvar, *request_handler_->dev_ctx(),
222
                            &reply_);
223
    }
X
Xin Pan 已提交
224
    Finish(reply_, &responder_);
G
gongweibao 已提交
225 226 227 228
  }

 protected:
  sendrecv::VariableMessage request_;
X
Xin Pan 已提交
229
  ::grpc::ByteBuffer reply_;
230
  ServerAsyncResponseWriter<::grpc::ByteBuffer> responder_;
G
gongweibao 已提交
231 232
};

233 234 235 236 237 238 239 240 241 242 243 244 245 246 247 248 249 250 251 252 253 254 255 256 257 258 259 260 261 262 263 264 265 266 267 268 269 270 271 272 273 274 275 276 277 278 279 280 281 282 283 284 285 286 287 288 289 290 291 292 293 294 295 296 297 298 299 300 301 302 303 304 305 306 307 308 309 310 311 312 313 314 315 316 317 318 319 320 321 322 323 324
class RequestGetMonomerVariable final : public RequestBase {
 public:
  explicit RequestGetMonomerVariable(GrpcService::AsyncService* service,
                                     ::grpc::ServerCompletionQueue* cq,
                                     RequestHandler* request_handler,
                                     int req_id, RPCServer* rpc_server)
      : RequestBase(service, cq, request_handler, req_id),
        responder_(&ctx_),
        rpc_server_(rpc_server) {
    auto method_id =
        static_cast<int>(distributed::GrpcMethod::kGetMonomerVariable);
    service_->RequestAsyncUnary(
        method_id, &ctx_, &request_, &responder_, cq_, cq_,
        reinterpret_cast<void*>(static_cast<intptr_t>(req_id)));
  }

  virtual ~RequestGetMonomerVariable() {}

  std::string GetReqName() override { return request_.varname(); }

  void Process() override {
    // proc request.
    std::string varname = request_.varname();

    rpc_server_->WaitVarCond(varname);
    MonomerHandle h = rpc_server_->GetMonomer(varname);

    auto scope = h.scope_;
    auto invar = scope->FindVar(varname);
    framework::Variable* outvar = nullptr;

    request_handler_->Handle(varname, scope, invar, &outvar,
                             request_.trainer_id());

    if (outvar) {
      SerializeToByteBuffer(varname, outvar, *h.dev_ctx_, &reply_);
    }
    Finish(reply_, &responder_);
  }

 protected:
  sendrecv::VariableMessage request_;
  ::grpc::ByteBuffer reply_;
  ServerAsyncResponseWriter<::grpc::ByteBuffer> responder_;
  RPCServer* rpc_server_{nullptr};
};

class RequestGetMonomerBarrier final : public RequestBase {
 public:
  explicit RequestGetMonomerBarrier(GrpcService::AsyncService* service,
                                    ::grpc::ServerCompletionQueue* cq,
                                    RequestHandler* request_handler, int req_id,
                                    RPCServer* rpc_server)
      : RequestBase(service, cq, request_handler, req_id),
        responder_(&ctx_),
        rpc_server_(rpc_server) {
    auto method_id =
        static_cast<int>(distributed::GrpcMethod::kGetMonomerBarrier);
    service_->RequestAsyncUnary(
        method_id, &ctx_, &request_, &responder_, cq_, cq_,
        reinterpret_cast<void*>(static_cast<intptr_t>(req_id)));
  }

  virtual ~RequestGetMonomerBarrier() {}

  std::string GetReqName() override { return request_.varname(); }

  void Process() override {
    // proc request.
    std::string varname = request_.varname();
    VLOG(4) << "RequestGetMonomerBarrier " << varname;

    rpc_server_->WaitVarCond(varname);
    MonomerHandle h = rpc_server_->GetMonomer(varname);

    framework::Scope* scope = nullptr;
    framework::Variable* invar = nullptr;
    framework::Variable* outvar = nullptr;

    request_handler_->Handle(varname, scope, invar, &outvar,
                             request_.trainer_id());

    Finish(reply_, &responder_);
  }

 protected:
  sendrecv::VariableMessage request_;
  sendrecv::VoidMessage reply_;
  ServerAsyncResponseWriter<sendrecv::VoidMessage> responder_;
  RPCServer* rpc_server_{nullptr};
};

325 326 327
class RequestPrefetch final : public RequestBase {
 public:
  explicit RequestPrefetch(GrpcService::AsyncService* service,
328 329 330
                           ::grpc::ServerCompletionQueue* cq,
                           RequestHandler* request_handler, int req_id)
      : RequestBase(service, cq, request_handler, req_id),
331
        responder_(&ctx_),
332
        local_scope_(nullptr) {
333 334
    request_.reset(new GRPCVariableResponse(request_handler->scope(),
                                            request_handler->dev_ctx(), true));
335 336
    int method_id =
        static_cast<int>(distributed::GrpcMethod::kPrefetchVariable);
X
Xin Pan 已提交
337 338
    service_->RequestAsyncUnary(
        method_id, &ctx_, request_.get(), &responder_, cq_, cq_,
339
        reinterpret_cast<void*>(static_cast<intptr_t>(req_id)));
340 341 342 343
  }

  virtual ~RequestPrefetch() {}

344
  std::string GetReqName() override { return request_->Varname(); }
345

346
  void Process() override {
347
    // prefetch process...
348 349
    std::string in_var_name = request_->Varname();
    std::string out_var_name = request_->OutVarname();
Q
Qiao Longfei 已提交
350
    std::string table_name = request_->TableName();
W
Wu Yi 已提交
351
    int trainer_id = request_->GetTrainerId();
352

M
minqiyang 已提交
353
    VLOG(4) << "RequestPrefetch, in_var_name: " << in_var_name
354
            << " out_var_name: " << out_var_name << " trainer: " << trainer_id;
355 356

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

W
Wu Yi 已提交
361
    request_handler_->Handle(in_var_name, scope, invar, &outvar, trainer_id,
Q
can run  
Qiao Longfei 已提交
362
                             out_var_name, table_name);
Y
Yancey1989 已提交
363

364
    SerializeToByteBuffer(out_var_name, outvar, *request_handler_->dev_ctx(),
365
                          &reply_);
X
Xin Pan 已提交
366
    Finish(reply_, &responder_);
367 368 369
  }

 protected:
370
  std::shared_ptr<GRPCVariableResponse> request_;
X
Xin Pan 已提交
371
  ::grpc::ByteBuffer reply_;
372
  ServerAsyncResponseWriter<::grpc::ByteBuffer> responder_;
373
  framework::Scope* local_scope_;
374 375
};

T
tangwei12 已提交
376 377 378 379 380
class RequestCheckpointNotify final : public RequestBase {
 public:
  explicit RequestCheckpointNotify(GrpcService::AsyncService* service,
                                   ::grpc::ServerCompletionQueue* cq,
                                   RequestHandler* request_handler, int req_id)
T
tangwei12 已提交
381
      : RequestBase(service, cq, request_handler, req_id), responder_(&ctx_) {
382 383
    request_.reset(new GRPCVariableResponse(request_handler->scope(),
                                            request_handler->dev_ctx()));
T
tangwei12 已提交
384 385
    int method_id =
        static_cast<int>(distributed::GrpcMethod::kCheckpointNotify);
T
tangwei12 已提交
386 387 388 389 390 391 392
    service_->RequestAsyncUnary(
        method_id, &ctx_, request_.get(), &responder_, cq_, cq_,
        reinterpret_cast<void*>(static_cast<intptr_t>(req_id)));
  }

  virtual ~RequestCheckpointNotify() {}

393
  std::string GetReqName() override { return request_->Varname(); }
T
tangwei12 已提交
394 395 396

  void Process() override {
    auto scope = request_->GetMutableLocalScope();
397 398

    std::string checkpoint_notify = request_->Varname();
T
tangwei12 已提交
399
    std::string checkpoint_dir = request_->OutVarname();
W
Wu Yi 已提交
400
    int trainer_id = request_->GetTrainerId();
401

M
minqiyang 已提交
402 403
    VLOG(4) << "RequestCheckpointNotify notify: " << checkpoint_notify
            << ", dir: " << checkpoint_dir;
T
tangwei12 已提交
404

T
tangwei12 已提交
405
    request_handler_->Handle(checkpoint_notify, scope, nullptr, nullptr,
W
Wu Yi 已提交
406
                             trainer_id, checkpoint_dir);
T
tangwei12 已提交
407 408
    Finish(reply_, &responder_);
  }
T
tangwei12 已提交
409 410

 protected:
411
  std::shared_ptr<GRPCVariableResponse> request_;
T
tangwei12 已提交
412 413
  sendrecv::VoidMessage reply_;
  ServerAsyncResponseWriter<sendrecv::VoidMessage> responder_;
T
tangwei12 已提交
414
};
T
tangwei12 已提交
415

416 417 418 419 420 421
class RequestNotify final : public RequestBase {
 public:
  explicit RequestNotify(GrpcService::AsyncService* service,
                         ::grpc::ServerCompletionQueue* cq,
                         RequestHandler* request_handler, int req_id)
      : RequestBase(service, cq, request_handler, req_id), responder_(&ctx_) {
1
123malin 已提交
422 423 424
    request_.reset(new GRPCVariableResponse(
        request_handler->scope(), request_handler->dev_ctx(),
        request_handler->distributed_mode()));
425 426 427 428 429 430 431 432 433
    int method_id = static_cast<int>(distributed::GrpcMethod::kRequestNotify);
    service_->RequestAsyncUnary(
        method_id, &ctx_, request_.get(), &responder_, cq_, cq_,
        reinterpret_cast<void*>(static_cast<intptr_t>(req_id)));
  }
  virtual ~RequestNotify() {}
  std::string GetReqName() override { return request_->Varname(); }

  void Process() override {
1
123malin 已提交
434 435
    std::string varname = GetReqName();
    VLOG(4) << "RequestNotify var_name:" << varname;
436

1
123malin 已提交
437 438
    auto scope = request_->GetMutableLocalScope();
    auto invar = request_->GetVar();
439
    int trainer_id = request_->GetTrainerId();
1
123malin 已提交
440 441
    framework::Variable* outvar = nullptr;
    request_handler_->Handle(varname, scope, invar, &outvar, trainer_id);
442 443 444 445 446
    Finish(reply_, &responder_);
  }

 protected:
  sendrecv::VoidMessage reply_;
1
123malin 已提交
447
  std::shared_ptr<GRPCVariableResponse> request_;
448 449 450
  ServerAsyncResponseWriter<sendrecv::VoidMessage> responder_;
};

451 452 453 454 455 456 457 458 459 460 461 462 463 464 465 466 467 468 469 470 471 472 473 474 475 476 477 478 479 480 481 482 483 484 485 486 487 488 489 490 491 492 493 494 495
class RequestSendAndRecv final : public RequestBase {
 public:
  explicit RequestSendAndRecv(GrpcService::AsyncService* service,
                              ::grpc::ServerCompletionQueue* cq,
                              RequestHandler* request_handler, int req_id)
      : RequestBase(service, cq, request_handler, req_id), responder_(&ctx_) {
    request_.reset(new GRPCVariableResponse(
        request_handler->scope(), request_handler->dev_ctx(),
        request_handler->distributed_mode()));

    int method_id =
        static_cast<int>(distributed::GrpcMethod::kRequestSendAndRecv);

    service_->RequestAsyncUnary(
        method_id, &ctx_, request_.get(), &responder_, cq_, cq_,
        reinterpret_cast<void*>(static_cast<intptr_t>(req_id)));
  }

  virtual ~RequestSendAndRecv() {}
  std::string GetReqName() override { return request_->Varname(); }

  void Process() override {
    std::string in_var_name = request_->Varname();
    std::string out_var_name = request_->OutVarname();
    std::string table_name = request_->TableName();
    int trainer_id = request_->GetTrainerId();

    VLOG(4) << "RequestSendAndRecv, in_var_name: " << in_var_name
            << " out_var_name: " << out_var_name << " trainer: " << trainer_id;
    auto scope = request_->GetMutableLocalScope();
    auto invar = scope->FindVar(in_var_name);
    framework::Variable* outvar = nullptr;
    request_handler_->Handle(in_var_name, scope, invar, &outvar, trainer_id,
                             out_var_name, table_name);
    SerializeToByteBuffer(out_var_name, outvar, *request_handler_->dev_ctx(),
                          &reply_);
    Finish(reply_, &responder_);
  }

 protected:
  std::shared_ptr<GRPCVariableResponse> request_;
  ::grpc::ByteBuffer reply_;
  ServerAsyncResponseWriter<::grpc::ByteBuffer> responder_;
};

T
done  
typhoonzero 已提交
496
void AsyncGRPCServer::WaitServerReady() {
497
  VLOG(4) << "AsyncGRPCServer is waiting server ready";
T
update  
typhoonzero 已提交
498
  std::unique_lock<std::mutex> lock(this->mutex_ready_);
T
done  
typhoonzero 已提交
499
  condition_ready_.wait(lock, [=] { return this->ready_ == 1; });
M
minqiyang 已提交
500
  VLOG(4) << "AsyncGRPCServer WaitSeverReady";
T
update  
typhoonzero 已提交
501 502
}

503 504 505 506 507 508 509 510 511 512 513 514 515 516
// Define an option subclass in order to disable SO_REUSEPORT for the
// server socket.
// Come from:
// https://github.com/tensorflow/tensorflow/blob/master/tensorflow/core/distributed_runtime/rpc/grpc_server_lib.cc
class NoReusePortOption : public ::grpc::ServerBuilderOption {
 public:
  void UpdateArguments(::grpc::ChannelArguments* args) override {
    args->SetInt(GRPC_ARG_ALLOW_REUSEPORT, 0);
  }

  void UpdatePlugins(std::vector<std::unique_ptr<::grpc::ServerBuilderPlugin>>*
                         plugins) override {}
};

517
void AsyncGRPCServer::StartServer() {
518 519 520 521 522 523 524 525 526 527 528 529
  for (int i = 0; i < FLAGS_rpc_retry_bind_port; i++) {
    ::grpc::ServerBuilder builder;
    std::unique_ptr<GrpcService::AsyncService> service(
        new GrpcService::AsyncService());
    builder.AddListeningPort(bind_address_, ::grpc::InsecureServerCredentials(),
                             &selected_port_);

    builder.SetMaxSendMessageSize(std::numeric_limits<int>::max());
    builder.SetMaxReceiveMessageSize(std::numeric_limits<int>::max());
    if (FLAGS_rpc_disable_reuse_port) {
      builder.SetOption(
          std::unique_ptr<::grpc::ServerBuilderOption>(new NoReusePortOption));
530
      LOG(INFO) << "set FLAGS_rpc_disable_reuse_port";
531 532 533 534 535 536 537 538 539 540 541 542 543 544 545 546 547 548
    }
    builder.RegisterService(service.get());

    for (auto t : rpc_call_map_) {
      rpc_cq_[t.first].reset(builder.AddCompletionQueue().release());
    }

    server_ = builder.BuildAndStart();
    if (selected_port_ != 0) {
      LOG(INFO) << "Server listening on " << bind_address_
                << " successful, selected port: " << selected_port_;
      service_.reset(service.release());
      break;
    }

    LOG(WARNING) << "Server listening on " << bind_address_
                 << " failed, selected port: " << selected_port_
                 << ", retry after 3 seconds!";
G
gongweibao 已提交
549

550
    sleep(3);
551
  }
Y
Yancey 已提交
552

553 554
  PADDLE_ENFORCE_NE(selected_port_, 0, "can't bind to address:%s",
                    bind_address_);
G
gongweibao 已提交
555

556 557 558
  std::function<void(const std::string&, int)> f =
      std::bind(&AsyncGRPCServer::TryToRegisterNewOne, this,
                std::placeholders::_1, std::placeholders::_2);
X
Xin Pan 已提交
559

560 561 562 563 564
  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 已提交
565

566 567 568
    reqs.reserve(kRequestBufSize);

    for (int i = 0; i < kRequestBufSize; i++) {
M
minqiyang 已提交
569
      VLOG(6) << "TryToRegisterNewOne on RPC NAME: " << rpc_name << " I: " << i;
570 571 572 573 574 575
      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)));
M
minqiyang 已提交
576
      VLOG(4) << t.first << " creates threads!";
577
    }
X
Xin Pan 已提交
578
  }
579

T
wip  
typhoonzero 已提交
580 581 582 583 584
  {
    std::lock_guard<std::mutex> lock(this->mutex_ready_);
    ready_ = 1;
  }
  condition_ready_.notify_all();
585

G
gongweibao 已提交
586 587
  // wait server
  server_->Wait();
588 589 590 591 592

  for (auto& t : rpc_threads_) {
    auto& threads = t.second;
    for (size_t i = 0; i < threads.size(); ++i) {
      threads[i]->join();
M
minqiyang 已提交
593
      VLOG(4) << t.first << " threads ends!";
594
    }
X
Xin Pan 已提交
595
  }
G
gongweibao 已提交
596 597 598
}

void AsyncGRPCServer::ShutdownQueue() {
599 600
  for (auto& t : rpc_cq_) {
    t.second->Shutdown();
M
minqiyang 已提交
601
    VLOG(4) << t.first << " queue shutdown!";
602
  }
G
gongweibao 已提交
603 604
}

605 606
void AsyncGRPCServer::ShutDownImpl() {
  std::unique_lock<std::mutex> lock(cq_mutex_);
T
typhoonzero 已提交
607
  is_shut_down_ = true;
G
gongweibao 已提交
608
  ShutdownQueue();
609

M
minqiyang 已提交
610
  VLOG(4) << "server_ shutdown!";
T
typhoonzero 已提交
611
  server_->Shutdown();
G
gongweibao 已提交
612 613
}

614 615
void AsyncGRPCServer::TryToRegisterNewOne(const std::string& rpc_name,
                                          int req_id) {
G
gongweibao 已提交
616 617
  std::unique_lock<std::mutex> lock(cq_mutex_);
  if (is_shut_down_) {
M
minqiyang 已提交
618
    VLOG(4) << "shutdown, do not TryToRegisterNewSendOne";
G
gongweibao 已提交
619 620 621
    return;
  }

M
minqiyang 已提交
622 623
  VLOG(4) << "TryToRegisterNewOne on RPC NAME: " << rpc_name
          << " REQ ID: " << req_id;
T
tangwei12 已提交
624

625 626 627 628 629 630
  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) {
631
    b = new RequestSend(service_.get(), cq.get(), handler, req_id);
632
  } else if (rpc_name == kRequestGet) {
633
    b = new RequestGet(service_.get(), cq.get(), handler, req_id);
634 635

  } else if (rpc_name == kRequestGetNoBarrier) {
636
    b = new RequestGetNoBarrier(service_.get(), cq.get(), handler, req_id);
637
  } else if (rpc_name == kRequestGetMonomerVariable) {
638
    b = new RequestGetMonomerVariable(service_.get(), cq.get(), handler, req_id,
639 640
                                      this);
  } else if (rpc_name == kRequestGetMonomerBarrier) {
641
    b = new RequestGetMonomerBarrier(service_.get(), cq.get(), handler, req_id,
642
                                     this);
643
  } else if (rpc_name == kRequestPrefetch) {
644
    b = new RequestPrefetch(service_.get(), cq.get(), handler, req_id);
T
tangwei12 已提交
645
  } else if (rpc_name == kRequestCheckpoint) {
646
    b = new RequestCheckpointNotify(service_.get(), cq.get(), handler, req_id);
647
  } else if (rpc_name == kRequestNotify) {
648
    b = new RequestNotify(service_.get(), cq.get(), handler, req_id);
649 650
  } else if (rpc_name == kRequestSendAndRecv) {
    b = new RequestSendAndRecv(service_.get(), cq.get(), handler, req_id);
651
  } else {
Q
qiaolongfei 已提交
652
    PADDLE_ENFORCE(false, "not supported rpc");
G
gongweibao 已提交
653 654
  }

655
  reqs[req_id] = b;
656

657
  VLOG(4) << "TryToRegisterNewOne status:" << b->Status();
658 659
}

X
Xin Pan 已提交
660
void AsyncGRPCServer::HandleRequest(
661 662
    ::grpc::ServerCompletionQueue* cq, const std::string& rpc_name,
    std::function<void(const std::string&, int)> TryToRegisterNewOne) {
G
gongweibao 已提交
663 664
  void* tag = NULL;
  bool ok = false;
665

G
gongweibao 已提交
666
  while (true) {
M
minqiyang 已提交
667
    VLOG(4) << "HandleRequest " << rpc_name << " wait next";
G
gongweibao 已提交
668
    if (!cq->Next(&tag, &ok)) {
669
      VLOG(4) << "CompletionQueue " << rpc_name << " shutdown!";
G
gongweibao 已提交
670 671
      break;
    }
Q
qiaolongfei 已提交
672

673
    int req_id = static_cast<int>(reinterpret_cast<intptr_t>(tag));
M
minqiyang 已提交
674 675
    VLOG(4) << "HandleRequest " << rpc_name << ", req_id:" << req_id
            << " get next";
G
gongweibao 已提交
676

677
    auto& reqs = rpc_reqs_[rpc_name];
X
Xin Pan 已提交
678 679
    RequestBase* base = nullptr;
    {
680 681 682
      PADDLE_ENFORCE(req_id >= 0 && req_id < kRequestBufSize);
      std::unique_lock<std::mutex> lock(cq_mutex_);
      base = reqs[req_id];
X
Xin Pan 已提交
683
    }
684

M
minqiyang 已提交
685
    VLOG(3) << base->Status2String(rpc_name);
G
gongweibao 已提交
686

G
gongweibao 已提交
687 688 689 690
    // 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 已提交
691
    if (!ok) {
G
gongweibao 已提交
692 693
      VLOG(4) << "completion queue:" << rpc_name << " recv no regular event"
              << " context:" << base->Status2String(rpc_name);
694
      TryToRegisterNewOne(rpc_name, req_id);
G
gongweibao 已提交
695 696 697 698 699 700 701 702 703 704
      delete base;
      continue;
    }

    switch (base->Status()) {
      case PROCESS: {
        base->Process();
        break;
      }
      case FINISH: {
705
        TryToRegisterNewOne(rpc_name, req_id);
G
gongweibao 已提交
706 707 708 709 710 711 712 713
        delete base;
        break;
      }
      default: { assert(false); }
    }
  }
}

714
}  // namespace distributed
G
gongweibao 已提交
715 716
}  // namespace operators
}  // namespace paddle