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

Y
Yi Wang 已提交
15
#include "paddle/fluid/operators/detail/grpc_server.h"
G
gongweibao 已提交
16 17 18 19 20 21 22 23 24 25 26 27 28 29 30

using grpc::ServerAsyncResponseWriter;

namespace paddle {
namespace operators {
namespace detail {

enum CallStatus { PROCESS = 0, FINISH };

// reference:
// https://stackoverflow.com/questions/41732884/grpc-multiple-services-in-cpp-async-server
class RequestBase {
 public:
  explicit RequestBase(sendrecv::SendRecvService::AsyncService* service,
                       grpc::ServerCompletionQueue* cq)
G
gongweibao 已提交
31 32 33
      : service_(service), cq_(cq), status_(PROCESS) {
    PADDLE_ENFORCE(cq_);
  }
G
gongweibao 已提交
34 35 36 37 38
  virtual ~RequestBase() {}
  virtual void Process() { assert(false); }

  CallStatus Status() { return status_; }
  void SetStatus(CallStatus status) { status_ = status; }
T
typhoonzero 已提交
39 40 41 42
  virtual std::string GetReqName() {
    assert(false);
    return "";
  }
G
gongweibao 已提交
43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64

 protected:
  grpc::ServerContext ctx_;
  sendrecv::SendRecvService::AsyncService* service_;
  grpc::ServerCompletionQueue* cq_;
  CallStatus status_;
};

typedef std::pair<std::string, sendrecv::VariableMessage> MessageWithName;

class RequestSend final : public RequestBase {
 public:
  explicit RequestSend(sendrecv::SendRecvService::AsyncService* service,
                       grpc::ServerCompletionQueue* cq,
                       SimpleBlockQueue<MessageWithName>* queue)
      : RequestBase(service, cq), queue_(queue), responder_(&ctx_) {
    service_->RequestSendVariable(&ctx_, &request_, &responder_, cq_, cq_,
                                  this);
  }

  virtual ~RequestSend() {}

G
gongweibao 已提交
65 66
  virtual std::string GetReqName() { return request_.varname(); }

G
gongweibao 已提交
67 68 69 70 71
  virtual void Process() {
    MessageWithName msg_with_name =
        std::make_pair(request_.varname(), std::move(request_));
    queue_->Push(std::move(msg_with_name));
    responder_.Finish(reply_, grpc::Status::OK, this);
G
gongweibao 已提交
72
    status_ = FINISH;
G
gongweibao 已提交
73 74 75 76 77 78 79 80 81 82 83 84
  }

 protected:
  sendrecv::VariableMessage request_;
  sendrecv::VoidMessage reply_;
  SimpleBlockQueue<MessageWithName>* queue_;
  ServerAsyncResponseWriter<sendrecv::VoidMessage> responder_;
};

class RequestGet final : public RequestBase {
 public:
  explicit RequestGet(sendrecv::SendRecvService::AsyncService* service,
Y
Yancey1989 已提交
85
                      grpc::ServerCompletionQueue* cq, framework::Scope* scope,
T
typhoonzero 已提交
86
                      const platform::DeviceContext* dev_ctx,
87
                      SimpleBlockQueue<MessageWithName>* queue)
Y
Yancey1989 已提交
88 89 90
      : RequestBase(service, cq),
        responder_(&ctx_),
        scope_(scope),
T
typhoonzero 已提交
91 92
        dev_ctx_(dev_ctx),
        queue_(queue) {
G
gongweibao 已提交
93 94 95 96 97
    service_->RequestGetVariable(&ctx_, &request_, &responder_, cq_, cq_, this);
  }

  virtual ~RequestGet() {}

G
gongweibao 已提交
98 99
  virtual std::string GetReqName() { return request_.varname(); }

G
gongweibao 已提交
100 101 102 103
  virtual void Process() {
    // proc request.
    std::string var_name = request_.varname();
    auto* var = scope_->FindVar(var_name);
104 105 106
    if (var_name != FETCH_BARRIER_MESSAGE) {
      SerializeToMessage(var_name, var, *dev_ctx_, &reply_);
    }
G
gongweibao 已提交
107 108
    // TODO(gongwb): check var's info.
    responder_.Finish(reply_, grpc::Status::OK, this);
G
gongweibao 已提交
109
    status_ = FINISH;
110 111 112 113
    MessageWithName msg_with_name =
        //          request name    reply
        std::make_pair(var_name, std::move(reply_));
    queue_->Push(msg_with_name);
G
gongweibao 已提交
114 115 116 117 118 119 120
  }

 protected:
  sendrecv::VariableMessage request_;
  sendrecv::VariableMessage reply_;
  ServerAsyncResponseWriter<sendrecv::VariableMessage> responder_;
  framework::Scope* scope_;
Y
Yancey1989 已提交
121
  const platform::DeviceContext* dev_ctx_;
122
  SimpleBlockQueue<MessageWithName>* queue_;
G
gongweibao 已提交
123 124
};

T
typhoonzero 已提交
125
void AsyncGRPCServer::WaitClientGet(int count) {
126 127 128 129 130 131
  int fetch_barriers = 0;
  while (fetch_barriers < count) {
    auto msg = var_get_queue_.Pop();
    if (msg.first == FETCH_BARRIER_MESSAGE) {
      fetch_barriers++;
    }
T
typhoonzero 已提交
132 133 134
  }
}

G
gongweibao 已提交
135 136 137
void AsyncGRPCServer::RunSyncUpdate() {
  grpc::ServerBuilder builder;
  builder.AddListeningPort(address_, grpc::InsecureServerCredentials());
G
gongweibao 已提交
138 139
  builder.SetMaxSendMessageSize(std::numeric_limits<int>::max());
  builder.SetMaxReceiveMessageSize(std::numeric_limits<int>::max());
G
gongweibao 已提交
140 141 142 143
  builder.RegisterService(&service_);

  cq_send_ = builder.AddCompletionQueue();
  cq_get_ = builder.AddCompletionQueue();
Y
Yancey 已提交
144

G
gongweibao 已提交
145 146 147 148 149 150 151 152 153
  server_ = builder.BuildAndStart();
  LOG(INFO) << "Server listening on " << address_ << std::endl;

  std::function<void()> send_register =
      std::bind(&AsyncGRPCServer::TryToRegisterNewSendOne, this);
  std::function<void()> get_register =
      std::bind(&AsyncGRPCServer::TryToRegisterNewGetOne, this);

  t_send_.reset(
Y
Yancey 已提交
154
      new std::thread(std::bind(&AsyncGRPCServer::HandleRequest, this,
G
gongweibao 已提交
155 156 157
                                cq_send_.get(), "cq_send", send_register)));

  t_get_.reset(
Y
Yancey 已提交
158
      new std::thread(std::bind(&AsyncGRPCServer::HandleRequest, this,
G
gongweibao 已提交
159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186
                                cq_get_.get(), "cq_get", get_register)));

  // wait server
  server_->Wait();
  t_send_->join();
  t_get_->join();
}

void AsyncGRPCServer::ShutdownQueue() {
  std::unique_lock<std::mutex> lock(cq_mutex_);
  cq_send_->Shutdown();
  cq_get_->Shutdown();
  is_shut_down_ = true;
}

// This URL explains why shutdown is complicate:
void AsyncGRPCServer::ShutDown() {
  server_->Shutdown();
  ShutdownQueue();
}

void AsyncGRPCServer::TryToRegisterNewSendOne() {
  std::unique_lock<std::mutex> lock(cq_mutex_);
  if (is_shut_down_) {
    return;
  }
  RequestSend* send =
      new RequestSend(&service_, cq_send_.get(), &var_recv_queue_);
Y
Yancey 已提交
187
  VLOG(4) << "Create RequestSend status:" << send->Status();
G
gongweibao 已提交
188 189 190 191 192 193 194
}

void AsyncGRPCServer::TryToRegisterNewGetOne() {
  std::unique_lock<std::mutex> lock(cq_mutex_);
  if (is_shut_down_) {
    return;
  }
T
typhoonzero 已提交
195 196
  RequestGet* get = new RequestGet(&service_, cq_get_.get(), scope_, dev_ctx_,
                                   &var_get_queue_);
Y
Yancey 已提交
197
  VLOG(4) << "Create RequestGet status:" << get->Status();
G
gongweibao 已提交
198 199
}

Y
Yancey 已提交
200 201
// FIXME(typhoonzero): change cq_name to enum.
void AsyncGRPCServer::HandleRequest(grpc::ServerCompletionQueue* cq,
G
gongweibao 已提交
202 203 204 205 206 207 208 209 210 211 212 213
                                    std::string cq_name,
                                    std::function<void()> TryToRegisterNewOne) {
  TryToRegisterNewOne();

  void* tag = NULL;
  bool ok = false;
  while (true) {
    if (!cq->Next(&tag, &ok)) {
      LOG(INFO) << cq_name << " get CompletionQueue shutdown!";
      break;
    }

G
gongweibao 已提交
214
    PADDLE_ENFORCE(tag);
T
typhoonzero 已提交
215 216
    // FIXME(typhoonzero): de-couple the barriers with recv_op
    if (cq_name == "cq_get") WaitCond(1);
T
typhoonzero 已提交
217
    if (cq_name == "cq_send") WaitCond(0);
G
gongweibao 已提交
218 219

    RequestBase* base = (RequestBase*)tag;
G
gongweibao 已提交
220 221 222 223
    // 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 已提交
224
    if (!ok) {
G
gongweibao 已提交
225 226
      LOG(WARNING) << cq_name << " recv no regular event:argument name"
                   << base->GetReqName();
G
gongweibao 已提交
227 228 229 230 231 232 233 234 235 236 237 238 239 240 241 242 243 244 245 246 247 248
      TryToRegisterNewOne();
      delete base;
      continue;
    }

    switch (base->Status()) {
      case PROCESS: {
        VLOG(4) << cq_name << " status:" << base->Status();
        TryToRegisterNewOne();
        base->Process();
        break;
      }
      case FINISH: {
        VLOG(4) << cq_name << " status:" << base->Status();
        delete base;
        break;
      }
      default: { assert(false); }
    }
  }
}

T
typhoonzero 已提交
249 250 251 252
void AsyncGRPCServer::WaitCond(int cond) {
  std::unique_lock<std::mutex> lock(this->barrier_mutex_);
  barrier_condition_.wait(lock,
                          [=] { return this->barrier_cond_step_ == cond; });
G
gongweibao 已提交
253 254
}

T
typhoonzero 已提交
255
void AsyncGRPCServer::SetCond(int cond) {
G
gongweibao 已提交
256
  {
T
typhoonzero 已提交
257 258
    std::lock_guard<std::mutex> lock(this->barrier_mutex_);
    barrier_cond_step_ = cond;
G
gongweibao 已提交
259
  }
T
typhoonzero 已提交
260
  barrier_condition_.notify_all();
G
gongweibao 已提交
261 262 263 264 265
}

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