recv_op.cc 7.0 KB
Newer Older
L
Luo Tao 已提交
1
/* Copyright (c) 2016 PaddlePaddle Authors. All Rights Reserve.
武毅 已提交
2

L
Luo Tao 已提交
3 4 5
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
武毅 已提交
6

L
Luo Tao 已提交
7
    http://www.apache.org/licenses/LICENSE-2.0
武毅 已提交
8

L
Luo Tao 已提交
9 10 11 12 13
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. */
武毅 已提交
14 15 16 17 18 19 20 21 22 23 24 25

#include <stdint.h>
#include <sys/stat.h>
#include <ostream>
#include <thread>

#include <unistd.h>

#include "paddle/framework/executor.h"
#include "paddle/framework/framework.pb.h"
#include "paddle/framework/lod_tensor.h"
#include "paddle/framework/op_registry.h"
T
typhoonzero 已提交
26
#include "paddle/framework/proto_desc.h"
G
gongweibao 已提交
27 28
#include "paddle/operators/detail/grpc_server.h"
#include "paddle/operators/detail/sendrecvop_utils.h"
武毅 已提交
29
#include "paddle/operators/detail/simple_block_queue.h"
T
typhoonzero 已提交
30
#include "paddle/string/printf.h"
武毅 已提交
31

T
typhoonzero 已提交
32 33
#define LISTEN_TERMINATE_MESSAGE "TERMINATE@RECV"

武毅 已提交
34 35 36
namespace paddle {
namespace operators {

37 38
constexpr char kOptimizeBlock[] = "OptimizeBlock";

G
gongweibao 已提交
39 40 41 42 43
void RunServer(std::shared_ptr<detail::AsyncGRPCServer> service) {
  service->RunSyncUpdate();
  VLOG(4) << "RunServer thread end";
}

Y
Yancey 已提交
44 45 46 47 48 49 50 51
static void CreateTensorFromMessageType(framework::Variable *var,
                                        sendrecv::VarType var_type) {
  if (var_type == sendrecv::VarType::LOD_TENSOR) {
    var->GetMutable<framework::LoDTensor>();
  } else if (var_type == sendrecv::VarType::SELECTED_ROWS) {
    var->GetMutable<framework::SelectedRows>();
  } else {
    PADDLE_THROW(
52
        "VariableMessage type %d is not in "
Y
Yancey 已提交
53 54 55 56 57
        "[LoDTensor, SelectedRows]",
        var_type);
  }
}

武毅 已提交
58 59 60 61 62 63 64 65
class RecvOp : public framework::OperatorBase {
 public:
  RecvOp(const std::string &type, const framework::VariableNameMap &inputs,
         const framework::VariableNameMap &outputs,
         const framework::AttributeMap &attrs)
      : OperatorBase(type, inputs, outputs, attrs) {
    if (!rpc_service_) {
      std::string endpoint = Attr<std::string>("endpoint");
G
gongweibao 已提交
66 67
      rpc_service_.reset(new detail::AsyncGRPCServer(endpoint));
      server_thread_.reset(new std::thread(RunServer, rpc_service_));
武毅 已提交
68 69 70
    }
  }

T
typhoonzero 已提交
71
  void Stop() override {
Y
Yancey 已提交
72
    detail::MessageWithName term_msg;
T
typhoonzero 已提交
73 74
    term_msg.first = LISTEN_TERMINATE_MESSAGE;
    rpc_service_->Push(term_msg);
G
gongweibao 已提交
75
    rpc_service_->ShutDown();
武毅 已提交
76 77 78
    server_thread_->join();
  }

T
done  
typhoonzero 已提交
79
  std::string GetGradVarNameForTrainer(const std::string &varname) const {
T
typhoonzero 已提交
80
    if (grads_counter_.find(varname) == grads_counter_.end()) {
T
done  
typhoonzero 已提交
81 82
      grads_counter_[varname] = 0;
    }
T
typhoonzero 已提交
83
    return string::Sprintf("%s.trainer_%d", varname, grads_counter_[varname]++);
T
done  
typhoonzero 已提交
84 85
  }

武毅 已提交
86
  void Run(const framework::Scope &scope,
D
dzhwinter 已提交
87
           const platform::Place &dev_place) const override {
Y
Yancey1989 已提交
88 89
    platform::DeviceContextPool &pool = platform::DeviceContextPool::Instance();
    auto &dev_ctx = *pool.Get(dev_place);
武毅 已提交
90
    framework::Scope &recv_scope = scope.NewScope();
Y
Yancey1989 已提交
91

Y
Yancey1989 已提交
92
    // FIXME(Yancey1989): initialize rpc server with laze mode.
T
typhoonzero 已提交
93
    rpc_service_->SetScope(&recv_scope);
Y
Yancey1989 已提交
94
    rpc_service_->SetDevCtx(&dev_ctx);
T
typhoonzero 已提交
95 96
    auto param_list = Attr<std::vector<std::string>>("ParamList");
    auto grad_list = Attr<std::vector<std::string>>("GradList");
T
typhoonzero 已提交
97
    auto fan_in = Attr<int>("Fanin");
T
typhoonzero 已提交
98
    size_t param_count = param_list.size();
G
gongweibao 已提交
99

100 101
    auto *block = Attr<framework::BlockDesc *>(kOptimizeBlock);
    auto *program = block->Program();
T
typhoonzero 已提交
102 103
    framework::Executor executor(dev_place);

T
typhoonzero 已提交
104
    // TODO(typhoonzero): change this to a while_op for every cluster-batch.
T
typhoonzero 已提交
105
    bool exit_flag = false;
106
    size_t barrier_size = param_count * fan_in;
T
typhoonzero 已提交
107
    while (!exit_flag) {
T
typhoonzero 已提交
108 109
      // Get from multiple trainers, we don't care about the order in which
      // the gradients arrives, just add suffix 0~n and merge the gradient.
T
typhoonzero 已提交
110 111
      rpc_service_->SetCond(0);
      for (size_t i = 0; i < barrier_size; ++i) {
Y
Yancey 已提交
112
        const detail::MessageWithName &v = rpc_service_->Get();
T
typhoonzero 已提交
113
        auto grad_var_name = v.first;
T
typhoonzero 已提交
114
        if (grad_var_name == LISTEN_TERMINATE_MESSAGE) {
T
typhoonzero 已提交
115
          LOG(INFO) << "received terminate message and exit";
T
typhoonzero 已提交
116 117 118
          exit_flag = true;
          break;
        }
T
typhoonzero 已提交
119 120 121 122
        auto it = std::find(grad_list.begin(), grad_list.end(), grad_var_name);
        std::string param_var_name;
        if (it != grad_list.end()) {
          param_var_name = param_list[it - grad_list.begin()];
T
typhoonzero 已提交
123
        } else {
124
          LOG(ERROR) << "grad has no paired param:" << grad_var_name;
T
typhoonzero 已提交
125
        }
126
        VLOG(3) << "received grad: " << grad_var_name
T
typhoonzero 已提交
127
                << " updating param: " << param_var_name;
T
typhoonzero 已提交
128
        if (fan_in > 1) {
T
typhoonzero 已提交
129
          grad_var_name = this->GetGradVarNameForTrainer(grad_var_name);
T
done  
typhoonzero 已提交
130
        }
131
        auto *var = recv_scope.FindVar(grad_var_name);
T
typhoonzero 已提交
132
        if (var == nullptr) {
133 134
          LOG(ERROR) << "Can not find server side var: " << grad_var_name;
          PADDLE_THROW("Can not find server side var");
T
done  
typhoonzero 已提交
135
        }
Y
Yancey 已提交
136
        detail::DeserializeFromMessage(v.second, dev_ctx, var);
T
typhoonzero 已提交
137 138 139
      }
      if (exit_flag) {
        break;
T
typhoonzero 已提交
140
      }
G
gongweibao 已提交
141

T
typhoonzero 已提交
142
      try {
143
        executor.Run(*program, &recv_scope, block->ID(), /*global_block*/
T
typhoonzero 已提交
144 145 146 147
                     false /*create_local_scope*/, false /*create_vars*/);
      } catch (std::exception &e) {
        LOG(ERROR) << "run sub program error " << e.what();
      }
T
typhoonzero 已提交
148 149
      rpc_service_->SetCond(1);
      rpc_service_->WaitClientGet(barrier_size);
T
typhoonzero 已提交
150
      grads_counter_.clear();
T
typhoonzero 已提交
151
    }  // while(true)
武毅 已提交
152 153 154
  }

 protected:
G
gongweibao 已提交
155
  std::shared_ptr<detail::AsyncGRPCServer> rpc_service_;
武毅 已提交
156
  std::shared_ptr<std::thread> server_thread_;
T
done  
typhoonzero 已提交
157
  mutable std::unordered_map<std::string, int> grads_counter_;
武毅 已提交
158 159 160 161
};

class RecvOpMaker : public framework::OpProtoAndCheckerMaker {
 public:
162
  RecvOpMaker(OpProto *proto, OpAttrChecker *op_checker)
武毅 已提交
163
      : OpProtoAndCheckerMaker(proto, op_checker) {
T
typhoonzero 已提交
164
    AddInput("RX", "(Tensor) Input tensor to be optimized").AsDuplicable();
武毅 已提交
165 166 167
    AddComment(R"DOC(
Recv operator

168
This operator will recieve tensor from send_op
武毅 已提交
169 170 171 172 173 174
)DOC");
    AddAttr<std::string>("endpoint",
                         "(string, default 127.0.0.1:6164)"
                         "IP address to listen on.")
        .SetDefault("127.0.0.1:6164")
        .AddCustomChecker([](const std::string &ip) { return !ip.empty(); });
175 176
    AddAttr<framework::BlockDesc *>(
        kOptimizeBlock, "Serialized ProgramDesc string for recv to run.");
T
typhoonzero 已提交
177 178
    AddAttr<std::vector<std::string>>(
        "ParamList", "type list of string",
179
        "grad->param name mapping to find which parameters to optimize.")
Y
Yancey1989 已提交
180
        .SetDefault({});
T
typhoonzero 已提交
181 182
    AddAttr<std::vector<std::string>>(
        "GradList", "type list of string",
183
        "grad->param name mapping to find which parameters to optimize.")
Y
Yancey1989 已提交
184
        .SetDefault({});
T
typhoonzero 已提交
185
    AddAttr<int>("Fanin", "type int",
T
done  
typhoonzero 已提交
186 187
                 "Number of trainers in the current cluster job")
        .SetDefault(1);
武毅 已提交
188 189 190 191 192 193 194 195 196
  }
};

}  // namespace operators
}  // namespace paddle

namespace ops = paddle::operators;

REGISTER_OPERATOR(recv, ops::RecvOp, ops::RecvOpMaker);