rpc_client.h 4.3 KB
Newer Older
G
gongweibao 已提交
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16
// Copyright (c) 2018 PaddlePaddle Authors. All Rights Reserved.
//
// 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.

#pragma once

17
#include <condition_variable>  // NOLINT
G
gongweibao 已提交
18
#include <string>
W
Wu Yi 已提交
19
#include "gflags/gflags.h"
G
gongweibao 已提交
20 21 22 23

#include "paddle/fluid/framework/data_type.h"
#include "paddle/fluid/framework/lod_tensor.h"
#include "paddle/fluid/framework/scope.h"
24
#include "paddle/fluid/operators/distributed/request_handler.h"
G
gongweibao 已提交
25

G
gongweibao 已提交
26
DECLARE_int32(rpc_deadline);
W
Wu Yi 已提交
27

G
gongweibao 已提交
28 29
namespace paddle {
namespace operators {
30
namespace distributed {
G
gongweibao 已提交
31 32 33

class RPCClient {
 public:
G
gongweibao 已提交
34 35
  RPCClient() {}
  virtual ~RPCClient() {}
36 37 38 39 40 41 42 43 44 45
  virtual VarHandlePtr AsyncSendVar(const std::string& ep,
                                    const platform::DeviceContext& ctx,
                                    const framework::Scope& scope,
                                    const std::string& var_name,
                                    int64_t time_out = FLAGS_rpc_deadline) = 0;

  virtual VarHandlePtr AsyncGetVar(const std::string& ep,
                                   const platform::DeviceContext& ctx,
                                   const framework::Scope& scope,
                                   const std::string& var_name,
46
                                   const std::string& out_varname,
Q
Qiao Longfei 已提交
47
                                   const std::string& table_name = "",
48 49
                                   int64_t time_out = FLAGS_rpc_deadline) = 0;

50 51 52 53 54 55
  virtual VarHandlePtr AsyncGetVarNoBarrier(
      const std::string& ep, const platform::DeviceContext& ctx,
      const framework::Scope& scope, const std::string& var_name,
      const std::string& out_varname,
      int64_t time_out = FLAGS_rpc_deadline) = 0;

56 57 58 59 60
  virtual VarHandlePtr AsyncGetMonomerVariable(
      const std::string& ep, const platform::DeviceContext& ctx,
      const framework::Scope& scope, const std::string& var_name,
      int64_t time_out = FLAGS_rpc_deadline) = 0;

61 62 63
  virtual VarHandlePtr AsyncPrefetchVar(
      const std::string& ep, const platform::DeviceContext& ctx,
      const framework::Scope& scope, const std::string& in_var_name,
Q
Qiao Longfei 已提交
64
      const std::string& out_var_name, const std::string& table_name = "",
65 66 67 68 69 70 71 72
      int64_t time_out = FLAGS_rpc_deadline) = 0;

  virtual VarHandlePtr AsyncSendBatchBarrier(
      const std::string& ep, int64_t time_out = FLAGS_rpc_deadline) = 0;

  virtual VarHandlePtr AsyncSendFetchBarrier(
      const std::string& ep, int64_t time_out = FLAGS_rpc_deadline) = 0;

73 74 75 76
  virtual VarHandlePtr AsyncGetMonomerBarrier(
      const std::string& ep, const std::string& var_name,
      int64_t time_out = FLAGS_rpc_deadline) = 0;

77 78 79 80 81 82
  virtual VarHandlePtr AsyncCheckpointNotify(
      const std::string& ep, const std::string& dir,
      int64_t time_out = FLAGS_rpc_deadline) = 0;

  virtual VarHandlePtr AsyncSendComplete(
      const std::string& ep, int64_t time_out = FLAGS_rpc_deadline) = 0;
Y
Yancey1989 已提交
83

Y
Yancey1989 已提交
84 85
  // Complete tells all the pserver instances that finishe the training,
  // the pserver can reduce it's barrier count, and continue to train
Y
Yancey1989 已提交
86
  // with other trainers.
Y
Yancey1989 已提交
87
  virtual void SendComplete() = 0;
W
Wu Yi 已提交
88

Y
Yancey1989 已提交
89
  virtual bool Wait() = 0;
G
gongweibao 已提交
90 91

  template <typename T>
W
Wu Yi 已提交
92 93
  static RPCClient* GetInstance(int trainer_id) {
    std::call_once(init_flag_, &RPCClient::Init<T>, trainer_id);
G
gongweibao 已提交
94 95 96 97 98
    return rpc_client_.get();
  }

  // Init is called by GetInstance.
  template <typename T>
W
Wu Yi 已提交
99
  static void Init(int trainer_id) {
Q
Qiao Longfei 已提交
100
    VLOG(0) << "init rpc client with trainer_id " << trainer_id;
W
Wu Yi 已提交
101
    trainer_id_ = trainer_id;
G
gongweibao 已提交
102 103 104 105 106 107 108
    if (rpc_client_.get() == nullptr) {
      rpc_client_.reset(new T());
      rpc_client_->InitImpl();
    }
  }

  virtual void InitImpl() {}
109 110

 protected:
W
Wu Yi 已提交
111 112
  // each trainer have exact one trainer id, it should be static
  static int trainer_id_;
G
gongweibao 已提交
113 114 115 116 117

 private:
  static std::once_flag init_flag_;
  static std::unique_ptr<RPCClient> rpc_client_;
};
118
}  // namespace distributed
G
gongweibao 已提交
119 120
}  // namespace operators
}  // namespace paddle