grpc_server.cc 24.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. */

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) {
M
MRXLT 已提交
60 61
    PADDLE_ENFORCE_NOT_NULL(cq_, platform::errors::InvalidArgument(
                                     "ServerCompletionQueue cq are empty"));
G
gongweibao 已提交
62
  }
G
gongweibao 已提交
63
  virtual ~RequestBase() {}
64
  virtual void Process() = 0;
G
gongweibao 已提交
65

G
gongweibao 已提交
66 67 68 69 70 71 72 73 74 75 76 77 78
  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 已提交
79 80 81 82 83 84 85 86 87 88 89 90
  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_)));
  }
91
  virtual std::string GetReqName() = 0;
G
gongweibao 已提交
92 93

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

class RequestSend final : public RequestBase {
 public:
105
  explicit RequestSend(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 110
    request_.reset(new GRPCVariableResponse(request_handler->scope(),
                                            request_handler->dev_ctx(), true));
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_) {
422 423
    request_.reset(new GRPCVariableResponse(request_handler->scope(),
                                            request_handler->dev_ctx(), true));
424 425 426 427 428 429 430 431 432
    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 已提交
433 434
    std::string varname = GetReqName();
    VLOG(4) << "RequestNotify var_name:" << varname;
435

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

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

450 451 452 453 454 455
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_) {
456 457
    request_.reset(new GRPCVariableResponse(request_handler->scope(),
                                            request_handler->dev_ctx(), true));
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

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

501 502 503 504 505 506 507 508 509 510 511 512 513 514
// 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 {}
};

515
void AsyncGRPCServer::StartServer() {
516 517 518 519 520 521 522 523 524 525 526 527
  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));
528
      LOG(INFO) << "set FLAGS_rpc_disable_reuse_port";
529 530 531 532 533 534 535 536 537 538 539 540 541 542 543 544 545 546
    }
    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 已提交
547

548
    sleep(3);
549
  }
Y
Yancey 已提交
550

M
MRXLT 已提交
551 552 553
  PADDLE_ENFORCE_NE(
      selected_port_, 0,
      platform::errors::Unavailable("can't bind to address:%s", bind_address_));
G
gongweibao 已提交
554

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

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

565 566 567
    reqs.reserve(kRequestBufSize);

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

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

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

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

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

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

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

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

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

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

  } else if (rpc_name == kRequestGetNoBarrier) {
635
    b = new RequestGetNoBarrier(service_.get(), cq.get(), handler, req_id);
636
  } else if (rpc_name == kRequestGetMonomerVariable) {
637
    b = new RequestGetMonomerVariable(service_.get(), cq.get(), handler, req_id,
638 639
                                      this);
  } else if (rpc_name == kRequestGetMonomerBarrier) {
640
    b = new RequestGetMonomerBarrier(service_.get(), cq.get(), handler, req_id,
641
                                     this);
642
  } else if (rpc_name == kRequestPrefetch) {
643
    b = new RequestPrefetch(service_.get(), cq.get(), handler, req_id);
T
tangwei12 已提交
644
  } else if (rpc_name == kRequestCheckpoint) {
645
    b = new RequestCheckpointNotify(service_.get(), cq.get(), handler, req_id);
646
  } else if (rpc_name == kRequestNotify) {
647
    b = new RequestNotify(service_.get(), cq.get(), handler, req_id);
648 649
  } else if (rpc_name == kRequestSendAndRecv) {
    b = new RequestSendAndRecv(service_.get(), cq.get(), handler, req_id);
650
  } else {
M
MRXLT 已提交
651 652
    PADDLE_THROW(
        platform::errors::InvalidArgument("not supported rpc: %s", rpc_name));
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;
    {
M
MRXLT 已提交
680 681 682 683
      PADDLE_ENFORCE_EQ(
          (req_id >= 0 && req_id < kRequestBufSize), true,
          platform::errors::OutOfRange("request id: %s out of bounds: [0, %s)",
                                       req_id, kRequestBufSize));
684 685
      std::unique_lock<std::mutex> lock(cq_mutex_);
      base = reqs[req_id];
X
Xin Pan 已提交
686
    }
687

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

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

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

717
}  // namespace distributed
G
gongweibao 已提交
718 719
}  // namespace operators
}  // namespace paddle