send_op.cc 4.4 KB
Newer Older
1
/* Copyright (c) 2016 PaddlePaddle Authors. All Rights Reserved.
武毅 已提交
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
#include <future>  // NOLINT
武毅 已提交
16 17
#include <ostream>

18
#include "paddle/fluid/framework/blocking_queue.h"
Y
Yi Wang 已提交
19 20 21
#include "paddle/fluid/framework/data_type.h"
#include "paddle/fluid/framework/lod_tensor.h"
#include "paddle/fluid/framework/op_registry.h"
W
Wu Yi 已提交
22
#include "paddle/fluid/operators/distributed/distributed.h"
W
Wu Yi 已提交
23
#include "paddle/fluid/operators/distributed_ops/send_recv_util.h"
24
#include "paddle/fluid/platform/profiler.h"
武毅 已提交
25 26 27 28 29 30

namespace paddle {
namespace operators {

class SendOp : public framework::OperatorBase {
 public:
G
gongweibao 已提交
31 32 33 34
  SendOp(const std::string& type, const framework::VariableNameMap& inputs,
         const framework::VariableNameMap& outputs,
         const framework::AttributeMap& attrs)
      : OperatorBase(type, inputs, outputs, attrs) {}
T
typhoonzero 已提交
35

T
typhoonzero 已提交
36 37
  void RunImpl(const framework::Scope& scope,
               const platform::Place& place) const override {
T
typhoonzero 已提交
38
    auto ins = Inputs("X");
G
gongweibao 已提交
39

40 41
    std::vector<std::string> epmap = Attr<std::vector<std::string>>("epmap");
    int sync_send = Attr<int>("sync_mode");
Q
qiaolongfei 已提交
42

Y
Yancey1989 已提交
43 44
    platform::DeviceContextPool& pool = platform::DeviceContextPool::Instance();
    auto& ctx = *pool.Get(place);
T
typhoonzero 已提交
45

46
    distributed::RPCClient* rpc_client =
W
Wu Yi 已提交
47 48
        distributed::RPCClient::GetInstance<RPCCLIENT_T>(
            Attr<int>("trainer_id"));
T
typhoonzero 已提交
49

50
    std::vector<distributed::VarHandlePtr> rets;
G
gongweibao 已提交
51
    for (size_t i = 0; i < ins.size(); i++) {
52
      if (NeedSend(scope, ins[i])) {
M
minqiyang 已提交
53
        VLOG(3) << "sending " << ins[i] << " to " << epmap[i];
54
        rets.push_back(rpc_client->AsyncSendVar(epmap[i], ctx, scope, ins[i]));
55
      } else {
M
minqiyang 已提交
56
        VLOG(3) << "don't send no-initialied variable: " << ins[i];
57
      }
武毅 已提交
58
    }
59
    if (sync_send) {
60
      for (size_t i = 0; i < rets.size(); i++) {
61
        VLOG(7) << "before sync_send " << ins[i] << "from " << epmap[i];
62
        PADDLE_ENFORCE(rets[i]->Wait(), "internal error in RPCClient");
63
        VLOG(7) << "after sync_send " << ins[i] << "from " << epmap[i];
64
      }
武毅 已提交
65 66 67 68 69 70
    }
  }
};

class SendOpMaker : public framework::OpProtoAndCheckerMaker {
 public:
Y
Yu Yang 已提交
71
  void Make() {
72
    AddInput("X", "(Tensor, SelectedRows) Input variables to be sent")
T
typhoonzero 已提交
73
        .AsDuplicable();
74 75
    AddOutput("Out", "(Any) Dummy outputs, used for control dependency")
        .AsDuplicable();
武毅 已提交
76
    AddComment(R"DOC(
A
Abhinav Arora 已提交
77
Send operator
武毅 已提交
78

79
This operator will send variables to listen_and_serve op at the parameter server.
武毅 已提交
80
)DOC");
81 82 83 84
    AddAttr<int>("sync_mode",
                 "(int, default 0)"
                 "sync send or async send.")
        .SetDefault(0);
W
Wu Yi 已提交
85
    AddAttr<int>("trainer_id", "trainer id from 0 ~ worker_num.").SetDefault(0);
T
typhoonzero 已提交
86 87 88
    AddAttr<std::vector<std::string>>("epmap",
                                      "(string vector, default 127.0.0.1:6164)"
                                      "Server endpoints in the order of input "
T
typhoonzero 已提交
89
                                      "variables for mapping")
90
        .SetDefault({"127.0.0.1:6164"});
Q
Qiao Longfei 已提交
91 92 93 94 95
    AddAttr<std::vector<int64_t>>("sections",
                                  "(vector<int>) "
                                  "the length of each output along the "
                                  "specified axis.")
        .SetDefault(std::vector<int64_t>{});
Q
Qiao Longfei 已提交
96 97 98 99 100 101 102 103 104 105
    AddAttr<std::vector<std::string>>(
        "send_varnames",
        "(vector<string>) "
        "the splited output varnames to send to pserver")
        .SetDefault(std::vector<std::string>{});
    AddAttr<int>("num",
                 "(int, default 0)"
                 "Number of sub-tensors. This must evenly divide "
                 "Input.dims()[axis]")
        .SetDefault(0);
武毅 已提交
106 107 108
  }
};

T
typhoonzero 已提交
109 110 111 112 113
class SendOpShapeInference : public framework::InferShapeBase {
 public:
  void operator()(framework::InferShapeContext* ctx) const override {}
};

武毅 已提交
114 115 116 117 118
}  // namespace operators
}  // namespace paddle

namespace ops = paddle::operators;

T
typhoonzero 已提交
119
REGISTER_OPERATOR(send, ops::SendOp, paddle::framework::EmptyGradOpMaker,
Y
Yancey1989 已提交
120
                  ops::SendOpMaker, ops::SendOpShapeInference);