|
|
|
@ -17,6 +17,7 @@ limitations under the License. */
|
|
|
|
|
#include <string>
|
|
|
|
|
#include <thread> // NOLINT
|
|
|
|
|
#include <utility>
|
|
|
|
|
#include <vector>
|
|
|
|
|
|
|
|
|
|
#include "grpc++/grpc++.h"
|
|
|
|
|
#include "paddle/fluid/framework/blocking_queue.h"
|
|
|
|
@ -30,6 +31,7 @@ limitations under the License. */
|
|
|
|
|
#include "paddle/fluid/operators/detail/send_recv.grpc.pb.h"
|
|
|
|
|
#include "paddle/fluid/operators/detail/send_recv.pb.h"
|
|
|
|
|
#include "paddle/fluid/operators/detail/sendrecvop_utils.h"
|
|
|
|
|
#include "paddle/fluid/platform/profiler.h"
|
|
|
|
|
|
|
|
|
|
namespace paddle {
|
|
|
|
|
namespace operators {
|
|
|
|
@ -82,19 +84,25 @@ class AsyncGRPCServer final {
|
|
|
|
|
protected:
|
|
|
|
|
void HandleRequest(::grpc::ServerCompletionQueue *cq,
|
|
|
|
|
const std::string &cq_name,
|
|
|
|
|
std::function<void()> TryToRegisterNewOne);
|
|
|
|
|
void TryToRegisterNewSendOne();
|
|
|
|
|
void TryToRegisterNewGetOne();
|
|
|
|
|
void TryToRegisterNewPrefetchOne();
|
|
|
|
|
std::function<void(int)> TryToRegisterNewOne);
|
|
|
|
|
void TryToRegisterNewSendOne(int i);
|
|
|
|
|
void TryToRegisterNewGetOne(int i);
|
|
|
|
|
void TryToRegisterNewPrefetchOne(int i);
|
|
|
|
|
void ShutdownQueue();
|
|
|
|
|
|
|
|
|
|
private:
|
|
|
|
|
static const int kSendReqsBufSize = 100;
|
|
|
|
|
static const int kGetReqsBufSize = 100;
|
|
|
|
|
|
|
|
|
|
std::mutex cq_mutex_;
|
|
|
|
|
volatile bool is_shut_down_ = false;
|
|
|
|
|
std::unique_ptr<::grpc::ServerCompletionQueue> cq_send_;
|
|
|
|
|
std::unique_ptr<::grpc::ServerCompletionQueue> cq_get_;
|
|
|
|
|
std::unique_ptr<::grpc::ServerCompletionQueue> cq_prefetch_;
|
|
|
|
|
|
|
|
|
|
RequestBase *send_reqs_[kSendReqsBufSize];
|
|
|
|
|
RequestBase *get_reqs_[kGetReqsBufSize];
|
|
|
|
|
|
|
|
|
|
GrpcService::AsyncService service_;
|
|
|
|
|
std::unique_ptr<::grpc::Server> server_;
|
|
|
|
|
|
|
|
|
@ -113,8 +121,9 @@ class AsyncGRPCServer final {
|
|
|
|
|
mutable int barrier_cond_step_;
|
|
|
|
|
std::condition_variable barrier_condition_;
|
|
|
|
|
|
|
|
|
|
std::unique_ptr<std::thread> t_send_;
|
|
|
|
|
std::unique_ptr<std::thread> t_get_;
|
|
|
|
|
std::vector<std::unique_ptr<std::thread>> t_sends_;
|
|
|
|
|
std::vector<std::unique_ptr<std::thread>> t_gets_;
|
|
|
|
|
|
|
|
|
|
std::unique_ptr<std::thread> t_prefetch_;
|
|
|
|
|
|
|
|
|
|
std::unique_ptr<framework::ExecutorPrepareContext> prefetch_ctx_;
|
|
|
|
|