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_) {
1
123malin 已提交
109 110 111
    request_.reset(new GRPCVariableResponse(
        request_handler->scope(), request_handler->dev_ctx(),
        request_handler->distributed_mode()));
112
    int method_id = static_cast<int>(distributed::GrpcMethod::kSendVariable);
X
Xin Pan 已提交
113 114
    service_->RequestAsyncUnary(
        method_id, &ctx_, request_.get(), &responder_, cq_, cq_,
X
Xin Pan 已提交
115
        reinterpret_cast<void*>(static_cast<intptr_t>(req_id)));
G
gongweibao 已提交
116 117
  }
  virtual ~RequestSend() {}
118 119 120 121
  std::string GetReqName() override { return request_->Varname(); }

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

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

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

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

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

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

  virtual ~RequestGet() {}

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

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

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

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

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

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

 protected:
  sendrecv::VariableMessage request_;
  ::grpc::ByteBuffer reply_;
Q
Qiao Longfei 已提交
185
  std::unique_ptr<framework::Scope> tmp_scope_;
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 222
  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(),
223
                            &reply_);
224
    }
X
Xin Pan 已提交
225
    Finish(reply_, &responder_);
G
gongweibao 已提交
226 227 228 229
  }

 protected:
  sendrecv::VariableMessage request_;
X
Xin Pan 已提交
230
  ::grpc::ByteBuffer reply_;
231
  ServerAsyncResponseWriter<::grpc::ByteBuffer> responder_;
G
gongweibao 已提交
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 325
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};
};

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

  virtual ~RequestPrefetch() {}

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

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

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

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

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

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

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

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

  virtual ~RequestCheckpointNotify() {}

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

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

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

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

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

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

417 418 419 420 421 422
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 已提交
423 424 425
    request_.reset(new GRPCVariableResponse(
        request_handler->scope(), request_handler->dev_ctx(),
        request_handler->distributed_mode()));
426 427 428 429 430 431 432 433 434
    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 已提交
435 436
    std::string varname = GetReqName();
    VLOG(4) << "RequestNotify var_name:" << varname;
437

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

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

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

504 505 506 507 508 509 510 511 512 513 514 515 516 517
// 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 {}
};

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

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

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

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

562 563 564 565 566
  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 已提交
567

568 569 570
    reqs.reserve(kRequestBufSize);

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

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

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

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

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

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

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

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

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

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

  } else if (rpc_name == kRequestGetNoBarrier) {
638
    b = new RequestGetNoBarrier(service_.get(), cq.get(), handler, req_id);
639
  } else if (rpc_name == kRequestGetMonomerVariable) {
640
    b = new RequestGetMonomerVariable(service_.get(), cq.get(), handler, req_id,
641 642
                                      this);
  } else if (rpc_name == kRequestGetMonomerBarrier) {
643
    b = new RequestGetMonomerBarrier(service_.get(), cq.get(), handler, req_id,
644
                                     this);
645
  } else if (rpc_name == kRequestPrefetch) {
646
    b = new RequestPrefetch(service_.get(), cq.get(), handler, req_id);
T
tangwei12 已提交
647
  } else if (rpc_name == kRequestCheckpoint) {
648
    b = new RequestCheckpointNotify(service_.get(), cq.get(), handler, req_id);
649
  } else if (rpc_name == kRequestNotify) {
650
    b = new RequestNotify(service_.get(), cq.get(), handler, req_id);
651 652
  } else if (rpc_name == kRequestSendAndRecv) {
    b = new RequestSendAndRecv(service_.get(), cq.get(), handler, req_id);
653
  } else {
M
MRXLT 已提交
654 655
    PADDLE_THROW(
        platform::errors::InvalidArgument("not supported rpc: %s", rpc_name));
G
gongweibao 已提交
656 657
  }

658
  reqs[req_id] = b;
659

660
  VLOG(4) << "TryToRegisterNewOne status:" << b->Status();
661 662
}

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

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

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

680
    auto& reqs = rpc_reqs_[rpc_name];
X
Xin Pan 已提交
681 682
    RequestBase* base = nullptr;
    {
M
MRXLT 已提交
683 684 685 686
      PADDLE_ENFORCE_EQ(
          (req_id >= 0 && req_id < kRequestBufSize), true,
          platform::errors::OutOfRange("request id: %s out of bounds: [0, %s)",
                                       req_id, kRequestBufSize));
687 688
      std::unique_lock<std::mutex> lock(cq_mutex_);
      base = reqs[req_id];
X
Xin Pan 已提交
689
    }
690

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

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

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

720
}  // namespace distributed
G
gongweibao 已提交
721 722
}  // namespace operators
}  // namespace paddle