grpc_serde.cc 6.1 KB
Newer Older
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17
/* Copyright (c) 2016 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. */

#ifdef PADDLE_WITH_CUDA
#include <nccl.h>
#endif
G
gongweibao 已提交
18
#include <limits>
19 20 21 22 23
#include <thread>  // NOLINT

#include "google/protobuf/io/coded_stream.h"
#include "google/protobuf/io/zero_copy_stream.h"
#include "paddle/fluid/framework/data_type.h"
W
Wu Yi 已提交
24 25 26
#include "paddle/fluid/operators/distributed/grpc/grpc_bytebuffer_stream.h"
#include "paddle/fluid/operators/distributed/grpc/grpc_serde.h"
#include "paddle/fluid/operators/distributed/grpc/grpc_variable_response.h"
27 28
#include "paddle/fluid/operators/distributed/proto_encoder_helper.h"
#include "paddle/fluid/operators/distributed/sendrecvop_utils.h"
P
peizhilin 已提交
29
#include "paddle/fluid/platform/port.h"
30 31 32 33 34 35 36 37
#include "paddle/fluid/platform/profiler.h"

namespace paddle {
namespace operators {
namespace distributed {

void SerializeToByteBuffer(const std::string& name, framework::Variable* var,
                           const platform::DeviceContext& ctx,
W
Wu Yi 已提交
38
                           ::grpc::ByteBuffer* msg, const std::string& out_name,
Q
Qiao Longfei 已提交
39 40
                           const int trainer_id,
                           const std::string& table_name) {
G
gongweibao 已提交
41
  platform::RecordRPCEvent record_event("serial", &ctx);
42
  VarMsg request;
Y
Yu Yang 已提交
43
  TensorPayload* payload = nullptr;
44 45

  request.set_varname(name);
W
Wu Yi 已提交
46
  request.set_trainer_id(trainer_id);
47 48 49 50 51 52 53 54 55 56 57 58 59 60
  // Note: normally the profiler is enabled in 1 trainer, hence only
  // 1 trainer returns true for ShouldSendProfileState(). It tells PS
  // servers the trainer's profiling state so that PS can follow the
  // trainer.
  if (platform::ShouldSendProfileState()) {
    if (platform::IsProfileEnabled()) {
      request.set_profile(platform::kEnableProfiler);
    } else {
      request.set_profile(platform::kDisableProfiler);
    }
  }
  if (!out_name.empty()) {
    request.set_out_varname(out_name);
  }
Q
Qiao Longfei 已提交
61 62 63
  if (!table_name.empty()) {
    request.set_table_name(table_name);
  }
64 65
  if (var->IsType<framework::LoDTensor>()) {
    request.set_type(::sendrecv::LOD_TENSOR);
Y
Yu Yang 已提交
66
    payload = new TensorPayload(GetTensorPayload(var, ctx, &request));
67 68
  } else if (var->IsType<framework::SelectedRows>()) {
    request.set_type(::sendrecv::SELECTED_ROWS);
Y
Yu Yang 已提交
69
    payload = new TensorPayload(GetSelectedRowsPayload(var, ctx, &request));
70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101
#ifdef PADDLE_WITH_CUDA
  } else if (var->IsType<ncclUniqueId>()) {
    request.set_type(::sendrecv::NCCL_ID);
#endif
  } else {
    PADDLE_THROW("Serialize does not support type: %s",
                 typeid(var->Type()).name());
  }

  std::string header;
  request.AppendToString(&header);
  auto buffer = std::unique_ptr<char[]>(new char[1024]);
  void* buf = buffer.get();
  ProtoEncodeHelper e(static_cast<char*>(buf), 1024);
  e.WriteRawBytes(std::string(header.data(), header.size()));
// NCCLID is copied directly to the message, return bytebuffer
// with only one slice if serializing NCCLID.
#ifdef PADDLE_WITH_CUDA
  if (var->IsType<ncclUniqueId>()) {
    e.WriteVarlengthBeginning(VarMsg::kSerializedFieldNumber,
                              NCCL_UNIQUE_ID_BYTES);
    const ncclUniqueId& uid = var->Get<ncclUniqueId>();
    e.WriteRawBytes(std::string(uid.internal, NCCL_UNIQUE_ID_BYTES));

    // for serialize NCCL_ID
    ::grpc::Slice slices(e.size());
    memcpy(const_cast<uint8_t*>(slices.begin()), e.data(), e.size());
    ::grpc::ByteBuffer tmp(&slices, 1);
    msg->Swap(&tmp);
    return;
  }
#endif
Y
Yu Yang 已提交
102
  PADDLE_ENFORCE_NOT_NULL(payload);
103

Y
Yu Yang 已提交
104
  e.WriteVarlengthBeginning(VarMsg::kSerializedFieldNumber,
Y
Yu Yang 已提交
105
                            payload->memory_size());
G
gongweibao 已提交
106 107 108 109
  if (payload->memory_size() >= std::numeric_limits<int>::max()) {
    LOG(FATAL) << "AppendZeroCopy varname:" << name
               << ", vlen:" << payload->memory_size();
  }
110 111 112 113 114
  // steal reference of tensor data
  ::grpc::Slice slices[4];  // metadata, tensor, rows meta, rows
  int num_slices = 2;       // only SelectedRows have rows buffer
  slices[0] = ::grpc::Slice(e.size());
  memcpy(const_cast<uint8_t*>(slices[0].begin()), e.data(), e.size());
Y
Yu Yang 已提交
115 116 117 118
  slices[1] = ::grpc::Slice(
      grpc_slice_new_with_user_data(payload->ptr(), payload->memory_size(),
                                    SerializeDestroyCallback, payload),
      ::grpc::Slice::STEAL_REF);
119 120 121 122

  if (var->IsType<framework::SelectedRows>()) {
    auto* slr = var->GetMutable<framework::SelectedRows>();
    ProtoEncodeHelper e2(static_cast<char*>(buf), 128);
G
gongweibao 已提交
123 124

    PADDLE_ENFORCE(VectorElemName(slr->rows()) == typeid(int64_t).name());
Y
Yu Yang 已提交
125
    size_t rows_memory_size = slr->rows().size() * sizeof(int64_t);
G
gongweibao 已提交
126

127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148
    e2.WriteVarlengthBeginning(VarMsg::kRowsFieldNumber, rows_memory_size);
    slices[2] = ::grpc::Slice(e2.size());
    memcpy(const_cast<uint8_t*>(slices[2].begin()), e2.data(), e2.size());

    slices[3] = ::grpc::Slice(
        grpc_slice_new_with_user_data(
            const_cast<void*>(
                reinterpret_cast<const void*>(slr->rows().data())),
            rows_memory_size, [](void* backing) {},
            const_cast<char*>(
                reinterpret_cast<const char*>(slr->rows().data()))),
        ::grpc::Slice::STEAL_REF);
    num_slices = 4;
  }

  ::grpc::ByteBuffer tmp(&slices[0], num_slices);
  msg->Swap(&tmp);
}

void DeserializeFromByteBuffer(const ::grpc::ByteBuffer& msg,
                               const platform::DeviceContext& ctx,
                               const framework::Scope* scope,
W
Wu Yi 已提交
149
                               framework::Variable** var, int* trainer_id) {
G
gongweibao 已提交
150
  platform::RecordRPCEvent record_event("deserial", &ctx);
151 152 153
  operators::distributed::GRPCVariableResponse resp(scope, &ctx);
  PADDLE_ENFORCE(resp.Parse(msg) == 0, "parse bytebuffer to tensor error!");
  *var = resp.GetVar();
W
Wu Yi 已提交
154
  *trainer_id = resp.GetTrainerId();
155 156 157 158 159
}

}  // namespace distributed
}  // namespace operators
}  // namespace paddle