data_set.cc 66.3 KB
Newer Older
1 2 3 4 5 6 7 8 9 10 11 12 13 14
/* 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. */

15
#include "paddle/fluid/framework/data_set.h"
16

17
#include "google/protobuf/text_format.h"
W
wangzhen38 已提交
18 19 20
#if (defined PADDLE_WITH_DISTRIBUTE) && (defined PADDLE_WITH_PSCORE)
#include "paddle/fluid/distributed/index_dataset/index_sampler.h"
#endif
21
#include "paddle/fluid/framework/data_feed_factory.h"
W
wangzhen38 已提交
22
#include "paddle/fluid/framework/fleet/fleet_wrapper.h"
23
#include "paddle/fluid/framework/io/fs.h"
H
hutuxian 已提交
24
#include "paddle/fluid/platform/monitor.h"
25
#include "paddle/fluid/platform/timer.h"
26
#include "paddle/phi/core/flags.h"
27

Z
zhaocaibei123 已提交
28
#ifdef PADDLE_WITH_PSCORE
29
#include "paddle/fluid/distributed/ps/wrapper/fleet.h"
D
danleifeng 已提交
30
#include "paddle/fluid/framework/fleet/heter_ps/graph_gpu_wrapper.h"
Z
zhaocaibei123 已提交
31 32
#endif

D
dongdaxiang 已提交
33 34 35 36 37
#if defined _WIN32 || defined __APPLE__
#else
#define _LINUX
#endif

H
hutuxian 已提交
38
USE_INT_STAT(STAT_total_feasign_num_in_mem);
L
lxsbupt 已提交
39
USE_INT_STAT(STAT_epoch_finish);
40 41
PHI_DECLARE_bool(graph_get_neighbor_id);
PHI_DECLARE_int32(gpugraph_storage_mode);
D
danleifeng 已提交
42

43 44 45
namespace paddle {
namespace framework {

X
xjqbest 已提交
46
// constructor
47
template <typename T>
D
dongdaxiang 已提交
48
DatasetImpl<T>::DatasetImpl() {
J
jiaqi 已提交
49
  VLOG(3) << "DatasetImpl<T>::DatasetImpl() constructor";
D
dongdaxiang 已提交
50
  thread_num_ = 1;
51
  trainer_num_ = 1;
J
jiaqi 已提交
52
  channel_num_ = 1;
53
  file_idx_ = 0;
H
hutuxian 已提交
54
  total_fea_num_ = 0;
J
jiaqi 已提交
55
  cur_channel_ = 0;
56 57
  fleet_send_batch_size_ = 1024;
  fleet_send_sleep_seconds_ = 0;
58
  merge_by_insid_ = false;
59 60
  merge_by_sid_ = true;
  enable_pv_merge_ = false;
61
  merge_size_ = 2;
62 63
  parse_ins_id_ = false;
  parse_content_ = false;
64
  parse_logkey_ = false;
65
  preload_thread_num_ = 0;
66
  global_index_ = 0;
67 68
  shuffle_by_uid_ = false;
  parse_uid_ = false;
D
dongdaxiang 已提交
69
}
70

X
xjqbest 已提交
71
// set filelist, file_idx_ will reset to zero.
72 73
template <typename T>
void DatasetImpl<T>::SetFileList(const std::vector<std::string>& filelist) {
74
  VLOG(3) << "filelist size: " << filelist.size();
75
  filelist_ = filelist;
76
  file_idx_ = 0;
77 78
}

X
xjqbest 已提交
79
// set expect thread num. actually it may change
80 81
template <typename T>
void DatasetImpl<T>::SetThreadNum(int thread_num) {
82
  VLOG(3) << "SetThreadNum thread_num=" << thread_num;
83 84 85
  thread_num_ = thread_num;
}

X
xjqbest 已提交
86 87 88
// if you run distributed, and want to do global shuffle,
// set this before global shuffle.
// be sure you call CreateReaders before SetTrainerNum
89
template <typename T>
X
xujiaqi01 已提交
90 91
void DatasetImpl<T>::SetTrainerNum(int trainer_num) {
  trainer_num_ = trainer_num;
92 93
}

X
xjqbest 已提交
94 95 96 97 98 99 100 101
// if you run distributed, and want to do global shuffle,
// set this before global shuffle.
// be sure you call CreateReaders before SetFleetSendBatchSize
template <typename T>
void DatasetImpl<T>::SetFleetSendBatchSize(int64_t size) {
  fleet_send_batch_size_ = size;
}

102 103 104
template <typename T>
void DatasetImpl<T>::SetHdfsConfig(const std::string& fs_name,
                                   const std::string& fs_ugi) {
X
xjqbest 已提交
105 106
  fs_name_ = fs_name;
  fs_ugi_ = fs_ugi;
107
  std::string cmd = std::string("$HADOOP_HOME/bin/hadoop fs");
108 109
  cmd += " -D fs.default.name=" + fs_name;
  cmd += " -D hadoop.job.ugi=" + fs_ugi;
110
  cmd += " -Ddfs.client.block.write.retries=15 -Ddfs.rpc.timeout=500000";
111
  paddle::framework::dataset_hdfs_set_command(cmd);
X
xujiaqi01 已提交
112
}
113

114 115 116 117 118 119 120 121 122 123
template <typename T>
void DatasetImpl<T>::SetDownloadCmd(const std::string& download_cmd) {
  paddle::framework::set_download_command(download_cmd);
}

template <typename T>
std::string DatasetImpl<T>::GetDownloadCmd() {
  return paddle::framework::download_cmd();
}

124 125
template <typename T>
void DatasetImpl<T>::SetDataFeedDesc(const std::string& data_feed_desc_str) {
126 127
  google::protobuf::TextFormat::ParseFromString(data_feed_desc_str,
                                                &data_feed_desc_);
128 129
}

Y
yaoxuefeng 已提交
130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147
template <typename T>
std::vector<std::string> DatasetImpl<T>::GetSlots() {
  auto multi_slot_desc = data_feed_desc_.multi_slot_desc();
  use_slots_.clear();
  for (int i = 0; i < multi_slot_desc.slots_size(); ++i) {
    const auto& slot = multi_slot_desc.slots(i);
    if (slot.type() == "uint64" || slot.type() == "uint32") {
      use_slots_.push_back(slot.name());
    }
  }
  std::cout << "dataset use slots: ";
  for (auto s : use_slots_) {
    std::cout << s << " | ";
  }
  std::cout << " end " << std::endl;
  return use_slots_;
}

148
template <typename T>
J
jiaqi 已提交
149 150 151 152
void DatasetImpl<T>::SetChannelNum(int channel_num) {
  channel_num_ = channel_num;
}

153 154 155 156 157 158 159 160 161 162
template <typename T>
void DatasetImpl<T>::SetParseInsId(bool parse_ins_id) {
  parse_ins_id_ = parse_ins_id;
}

template <typename T>
void DatasetImpl<T>::SetParseContent(bool parse_content) {
  parse_content_ = parse_content;
}

163 164 165 166 167
template <typename T>
void DatasetImpl<T>::SetParseLogKey(bool parse_logkey) {
  parse_logkey_ = parse_logkey;
}

168
template <typename T>
169
void DatasetImpl<T>::SetMergeByInsId(int merge_size) {
170
  merge_by_insid_ = true;
171
  parse_ins_id_ = true;
172
  merge_size_ = merge_size;
173 174
}

175 176 177 178 179
template <typename T>
void DatasetImpl<T>::SetMergeBySid(bool is_merge) {
  merge_by_sid_ = is_merge;
}

180 181 182 183 184 185
template <typename T>
void DatasetImpl<T>::SetShuffleByUid(bool enable_shuffle_uid) {
  shuffle_by_uid_ = enable_shuffle_uid;
  parse_uid_ = true;
}

186 187 188 189 190
template <typename T>
void DatasetImpl<T>::SetEnablePvMerge(bool enable_pv_merge) {
  enable_pv_merge_ = enable_pv_merge;
}

191 192 193 194 195 196
template <typename T>
void DatasetImpl<T>::SetGenerateUniqueFeasign(bool gen_uni_feasigns) {
  gen_uni_feasigns_ = gen_uni_feasigns;
  VLOG(3) << "Set generate unique feasigns: " << gen_uni_feasigns;
}

197 198 199 200 201 202 203 204
template <typename T>
void DatasetImpl<T>::SetFeaEval(bool fea_eval, int record_candidate_size) {
  slots_shuffle_fea_eval_ = fea_eval;
  slots_shuffle_rclist_.ReSize(record_candidate_size);
  VLOG(3) << "SetFeaEval fea eval mode: " << fea_eval
          << " with record candidate size: " << record_candidate_size;
}

D
danleifeng 已提交
205 206 207 208 209 210 211 212 213 214
template <typename T>
void DatasetImpl<T>::SetGpuGraphMode(int is_graph_mode) {
  gpu_graph_mode_ = is_graph_mode;
}

template <typename T>
int DatasetImpl<T>::GetGpuGraphMode() {
  return gpu_graph_mode_;
}

J
jiaqi 已提交
215 216 217 218 219 220 221 222 223 224 225 226 227 228 229
template <typename T>
std::vector<paddle::framework::DataFeed*> DatasetImpl<T>::GetReaders() {
  std::vector<paddle::framework::DataFeed*> ret;
  ret.reserve(readers_.size());
  for (auto i : readers_) {
    ret.push_back(i.get());
  }
  return ret;
}

template <typename T>
void DatasetImpl<T>::CreateChannel() {
  if (input_channel_ == nullptr) {
    input_channel_ = paddle::framework::MakeChannel<T>();
  }
230
  if (multi_output_channel_.empty()) {
J
jiaqi 已提交
231 232 233 234 235
    multi_output_channel_.reserve(channel_num_);
    for (int i = 0; i < channel_num_; ++i) {
      multi_output_channel_.push_back(paddle::framework::MakeChannel<T>());
    }
  }
236
  if (multi_consume_channel_.empty()) {
J
jiaqi 已提交
237 238 239 240 241
    multi_consume_channel_.reserve(channel_num_);
    for (int i = 0; i < channel_num_; ++i) {
      multi_consume_channel_.push_back(paddle::framework::MakeChannel<T>());
    }
  }
242 243 244
  if (input_pv_channel_ == nullptr) {
    input_pv_channel_ = paddle::framework::MakeChannel<PvInstance>();
  }
245
  if (multi_pv_output_.empty()) {
246 247 248 249 250
    multi_pv_output_.reserve(channel_num_);
    for (int i = 0; i < channel_num_; ++i) {
      multi_pv_output_.push_back(paddle::framework::MakeChannel<PvInstance>());
    }
  }
251
  if (multi_pv_consume_.empty()) {
252 253 254 255 256
    multi_pv_consume_.reserve(channel_num_);
    for (int i = 0; i < channel_num_; ++i) {
      multi_pv_consume_.push_back(paddle::framework::MakeChannel<PvInstance>());
    }
  }
257 258
}

259 260 261
// if sent message between workers, should first call this function
template <typename T>
void DatasetImpl<T>::RegisterClientToClientMsgHandler() {
Z
zhaocaibei123 已提交
262 263 264 265 266 267
#ifdef PADDLE_WITH_PSCORE
  auto fleet_ptr = distributed::FleetWrapper::GetInstance();
#else
  auto fleet_ptr = framework::FleetWrapper::GetInstance();
#endif
  VLOG(1) << "RegisterClientToClientMsgHandler";
268 269 270 271
  fleet_ptr->RegisterClientToClientMsgHandler(
      0, [this](int msg_type, int client_id, const std::string& msg) -> int {
        return this->ReceiveFromClient(msg_type, client_id, msg);
      });
Z
zhaocaibei123 已提交
272
  VLOG(1) << "RegisterClientToClientMsgHandler done";
273
}
274 275
static void compute_left_batch_num(const int ins_num,
                                   const int thread_num,
Y
yaoxuefeng 已提交
276 277 278 279 280 281 282 283 284 285 286 287 288 289 290
                                   std::vector<std::pair<int, int>>* offset,
                                   const int start_pos) {
  int cur_pos = start_pos;
  int batch_size = ins_num / thread_num;
  int left_num = ins_num % thread_num;
  for (int i = 0; i < thread_num; ++i) {
    int batch_num_size = batch_size;
    if (i == 0) {
      batch_num_size = batch_num_size + left_num;
    }
    offset->push_back(std::make_pair(cur_pos, batch_num_size));
    cur_pos += batch_num_size;
  }
}

291 292
static void compute_batch_num(const int64_t ins_num,
                              const int batch_size,
Y
yaoxuefeng 已提交
293 294 295 296 297 298 299 300 301 302 303 304 305 306 307 308 309 310 311 312 313 314 315 316 317 318 319 320 321 322 323 324 325
                              const int thread_num,
                              std::vector<std::pair<int, int>>* offset) {
  int thread_batch_num = batch_size * thread_num;
  // less data
  if (static_cast<int64_t>(thread_batch_num) > ins_num) {
    compute_left_batch_num(ins_num, thread_num, offset, 0);
    return;
  }

  int cur_pos = 0;
  int offset_num = static_cast<int>(ins_num / thread_batch_num) * thread_num;
  int left_ins_num = static_cast<int>(ins_num % thread_batch_num);
  if (left_ins_num > 0 && left_ins_num < thread_num) {
    offset_num = offset_num - thread_num;
    left_ins_num = left_ins_num + thread_batch_num;
    for (int i = 0; i < offset_num; ++i) {
      offset->push_back(std::make_pair(cur_pos, batch_size));
      cur_pos += batch_size;
    }
    // split data to thread avg two rounds
    compute_left_batch_num(left_ins_num, thread_num * 2, offset, cur_pos);
  } else {
    for (int i = 0; i < offset_num; ++i) {
      offset->push_back(std::make_pair(cur_pos, batch_size));
      cur_pos += batch_size;
    }
    if (left_ins_num > 0) {
      compute_left_batch_num(left_ins_num, thread_num, offset, cur_pos);
    }
  }
}

static int compute_thread_batch_nccl(
326 327 328 329
    const int thr_num,
    const int64_t total_instance_num,
    const int minibatch_size,
    std::vector<std::pair<int, int>>* nccl_offsets) {
Y
yaoxuefeng 已提交
330 331 332 333 334 335 336 337 338 339 340 341 342 343 344
  int thread_avg_batch_num = 0;
  if (total_instance_num < static_cast<int64_t>(thr_num)) {
    LOG(WARNING) << "compute_thread_batch_nccl total ins num:["
                 << total_instance_num << "], less thread num:[" << thr_num
                 << "]";
    return thread_avg_batch_num;
  }

  auto& offset = (*nccl_offsets);
  // split data avg by thread num
  compute_batch_num(total_instance_num, minibatch_size, thr_num, &offset);
  thread_avg_batch_num = static_cast<int>(offset.size() / thr_num);
#ifdef PADDLE_WITH_GLOO
  auto gloo_wrapper = paddle::framework::GlooWrapper::GetInstance();
  if (gloo_wrapper->Size() > 1) {
345 346 347 348
    if (!gloo_wrapper->IsInitialized()) {
      VLOG(0) << "GLOO is not inited";
      gloo_wrapper->Init();
    }
Y
yaoxuefeng 已提交
349 350 351 352 353 354 355 356 357 358 359 360 361 362 363 364 365 366 367 368 369 370 371 372 373
    // adjust batch num per thread for NCCL
    std::vector<int> thread_avg_batch_num_vec(1, thread_avg_batch_num);
    std::vector<int64_t> total_instance_num_vec(1, total_instance_num);
    auto thread_max_batch_num_vec =
        gloo_wrapper->AllReduce(thread_avg_batch_num_vec, "max");
    auto sum_total_ins_num_vec =
        gloo_wrapper->AllReduce(total_instance_num_vec, "sum");
    int thread_max_batch_num = thread_max_batch_num_vec[0];
    int64_t sum_total_ins_num = sum_total_ins_num_vec[0];
    int diff_batch_num = thread_max_batch_num - thread_avg_batch_num;
    VLOG(3) << "diff batch num: " << diff_batch_num
            << " thread max batch num: " << thread_max_batch_num
            << " thread avg batch num: " << thread_avg_batch_num;
    if (diff_batch_num == 0) {
      LOG(WARNING) << "total sum ins " << sum_total_ins_num << ", thread_num "
                   << thr_num << ", ins num " << total_instance_num
                   << ", batch num " << offset.size()
                   << ", thread avg batch num " << thread_avg_batch_num;
      return thread_avg_batch_num;
    }

    int need_ins_num = thread_max_batch_num * thr_num;
    // data is too less
    if ((int64_t)need_ins_num > total_instance_num) {
      PADDLE_THROW(platform::errors::InvalidArgument(
374 375
          "error instance num:[%d] less need ins num:[%d]",
          total_instance_num,
Y
yaoxuefeng 已提交
376 377 378 379 380 381 382 383 384 385 386 387 388 389
          need_ins_num));
      return thread_avg_batch_num;
    }

    int need_batch_num = (diff_batch_num + 1) * thr_num;
    int offset_split_index = static_cast<int>(offset.size() - thr_num);
    int split_left_num = total_instance_num - offset[offset_split_index].first;
    while (split_left_num < need_batch_num) {
      need_batch_num += thr_num;
      offset_split_index -= thr_num;
      split_left_num = total_instance_num - offset[offset_split_index].first;
    }
    int split_start = offset[offset_split_index].first;
    offset.resize(offset_split_index);
390 391
    compute_left_batch_num(
        split_left_num, need_batch_num, &offset, split_start);
Y
yaoxuefeng 已提交
392 393 394 395 396 397 398 399 400 401 402 403 404 405 406 407 408 409 410 411 412
    LOG(WARNING) << "total sum ins " << sum_total_ins_num << ", thread_num "
                 << thr_num << ", ins num " << total_instance_num
                 << ", batch num " << offset.size() << ", thread avg batch num "
                 << thread_avg_batch_num << ", thread max batch num "
                 << thread_max_batch_num
                 << ", need batch num: " << (need_batch_num / thr_num)
                 << "split begin (" << split_start << ")" << split_start
                 << ", num " << split_left_num;
    thread_avg_batch_num = thread_max_batch_num;
  } else {
    LOG(WARNING) << "thread_num " << thr_num << ", ins num "
                 << total_instance_num << ", batch num " << offset.size()
                 << ", thread avg batch num " << thread_avg_batch_num;
  }
#else
  PADDLE_THROW(platform::errors::Unavailable(
      "dataset compute nccl batch number need compile with GLOO"));
#endif
  return thread_avg_batch_num;
}

Y
yaoxuefeng 已提交
413
void MultiSlotDataset::PrepareTrain() {
Y
yaoxuefeng 已提交
414 415
#ifdef PADDLE_WITH_GLOO
  if (enable_heterps_) {
416
    if (input_records_.empty() && input_channel_ != nullptr &&
Y
yaoxuefeng 已提交
417 418 419 420 421 422 423 424 425 426 427 428 429 430
        input_channel_->Size() != 0) {
      input_channel_->ReadAll(input_records_);
      VLOG(3) << "read from channel to records with records size: "
              << input_records_.size();
    }
    VLOG(3) << "input records size: " << input_records_.size();
    int64_t total_ins_num = input_records_.size();
    std::vector<std::pair<int, int>> offset;
    int default_batch_size =
        reinterpret_cast<MultiSlotInMemoryDataFeed*>(readers_[0].get())
            ->GetDefaultBatchSize();
    VLOG(3) << "thread_num: " << thread_num_
            << " memory size: " << total_ins_num
            << " default batch_size: " << default_batch_size;
431 432
    compute_thread_batch_nccl(
        thread_num_, total_ins_num, default_batch_size, &offset);
Y
yaoxuefeng 已提交
433 434 435 436 437 438 439 440 441 442 443 444 445 446 447 448 449
    VLOG(3) << "offset size: " << offset.size();
    for (int i = 0; i < thread_num_; i++) {
      reinterpret_cast<MultiSlotInMemoryDataFeed*>(readers_[i].get())
          ->SetRecord(&input_records_[0]);
    }
    for (size_t i = 0; i < offset.size(); i++) {
      reinterpret_cast<MultiSlotInMemoryDataFeed*>(
          readers_[i % thread_num_].get())
          ->AddBatchOffset(offset[i]);
    }
  }
#else
  PADDLE_THROW(platform::errors::Unavailable(
      "dataset set heterps need compile with GLOO"));
#endif
  return;
}
450

X
xjqbest 已提交
451 452
// load data into memory, Dataset hold this memory,
// which will later be fed into readers' channel
453 454 455
template <typename T>
void DatasetImpl<T>::LoadIntoMemory() {
  VLOG(3) << "DatasetImpl<T>::LoadIntoMemory() begin";
456 457
  platform::Timer timeline;
  timeline.Start();
458
  std::vector<std::thread> load_threads;
D
danleifeng 已提交
459
  if (gpu_graph_mode_) {
460
    VLOG(1) << "in gpu_graph_mode";
L
lxsbupt 已提交
461 462 463 464
#if defined(PADDLE_WITH_GPU_GRAPH) && defined(PADDLE_WITH_HETERPS)
    for (size_t i = 0; i < readers_.size(); i++) {
      readers_[i]->SetGpuGraphMode(gpu_graph_mode_);
    }
D
danleifeng 已提交
465

L
lxsbupt 已提交
466 467 468
    if (STAT_GET(STAT_epoch_finish) == 1) {
      VLOG(0) << "get epoch finish true";
      STAT_RESET(STAT_epoch_finish, 0);
D
danleifeng 已提交
469
      for (size_t i = 0; i < readers_.size(); i++) {
L
lxsbupt 已提交
470 471
        readers_[i]->ResetPathNum();
        readers_[i]->ResetEpochFinish();
D
danleifeng 已提交
472 473 474
      }
    }

L
lxsbupt 已提交
475 476 477 478 479 480 481 482 483 484 485 486 487 488 489 490 491
    for (int64_t i = 0; i < thread_num_; ++i) {
      load_threads.push_back(std::thread(
          &paddle::framework::DataFeed::DoWalkandSage, readers_[i].get()));
    }
    for (std::thread& t : load_threads) {
      t.join();
    }
    uint64_t node_num = 0;
    for (int i = 0; i < thread_num_; i++) {
      auto host_vec = readers_[i]->GetHostVec();
      node_num += host_vec->size();
    }
    gpu_graph_total_keys_.reserve(node_num);
    for (int i = 0; i < thread_num_; i++) {
      auto host_vec = readers_[i]->GetHostVec();
      for (size_t j = 0; j < host_vec->size(); j++) {
        gpu_graph_total_keys_.push_back((*host_vec)[j]);
D
danleifeng 已提交
492 493 494
      }
    }

L
lxsbupt 已提交
495 496 497 498 499
    if (GetEpochFinish() == true) {
      VLOG(0) << "epoch finish, set stat and clear sample stat!";
      STAT_RESET(STAT_epoch_finish, 1);
      for (size_t i = 0; i < readers_.size(); i++) {
        readers_[i]->ClearSampleState();
D
danleifeng 已提交
500
      }
L
lxsbupt 已提交
501 502 503 504
    }
    if (FLAGS_gpugraph_storage_mode != GpuGraphStorageMode::WHOLE_HBM) {
      for (size_t i = 0; i < readers_.size(); i++) {
        readers_[i]->clear_gpu_mem();
D
danleifeng 已提交
505 506
      }
    }
L
lxsbupt 已提交
507 508 509

    VLOG(2) << "end add edge into gpu_graph_total_keys_ size["
            << gpu_graph_total_keys_.size() << "]";
D
danleifeng 已提交
510 511 512
#endif
  } else {
    for (int64_t i = 0; i < thread_num_; ++i) {
513 514
      load_threads.emplace_back(&paddle::framework::DataFeed::LoadIntoMemory,
                                readers_[i].get());
D
danleifeng 已提交
515 516 517 518
    }
    for (std::thread& t : load_threads) {
      t.join();
    }
519
  }
J
jiaqi 已提交
520 521 522
  input_channel_->Close();
  int64_t in_chan_size = input_channel_->Size();
  input_channel_->SetBlockSize(in_chan_size / thread_num_ + 1);
523

524 525
  timeline.Pause();
  VLOG(3) << "DatasetImpl<T>::LoadIntoMemory() end"
J
jiaqi 已提交
526
          << ", memory data size=" << input_channel_->Size()
527
          << ", cost time=" << timeline.ElapsedSec() << " seconds";
528 529
}

J
jiaqi 已提交
530 531 532
template <typename T>
void DatasetImpl<T>::PreLoadIntoMemory() {
  VLOG(3) << "DatasetImpl<T>::PreLoadIntoMemory() begin";
533
  if (preload_thread_num_ != 0) {
534
    CHECK(static_cast<size_t>(preload_thread_num_) == preload_readers_.size());
535 536
    preload_threads_.clear();
    for (int64_t i = 0; i < preload_thread_num_; ++i) {
537 538 539
      preload_threads_.emplace_back(
          &paddle::framework::DataFeed::LoadIntoMemory,
          preload_readers_[i].get());
540 541
    }
  } else {
542
    CHECK(static_cast<size_t>(thread_num_) == readers_.size());
543 544
    preload_threads_.clear();
    for (int64_t i = 0; i < thread_num_; ++i) {
545 546
      preload_threads_.emplace_back(
          &paddle::framework::DataFeed::LoadIntoMemory, readers_[i].get());
547
    }
J
jiaqi 已提交
548 549 550 551 552 553 554 555 556 557 558 559 560 561 562 563
  }
  VLOG(3) << "DatasetImpl<T>::PreLoadIntoMemory() end";
}

template <typename T>
void DatasetImpl<T>::WaitPreLoadDone() {
  VLOG(3) << "DatasetImpl<T>::WaitPreLoadDone() begin";
  for (std::thread& t : preload_threads_) {
    t.join();
  }
  input_channel_->Close();
  int64_t in_chan_size = input_channel_->Size();
  input_channel_->SetBlockSize(in_chan_size / thread_num_ + 1);
  VLOG(3) << "DatasetImpl<T>::WaitPreLoadDone() end";
}

564 565 566
// release memory data
template <typename T>
void DatasetImpl<T>::ReleaseMemory() {
T
Thunderbrook 已提交
567 568 569 570 571
  release_thread_ = new std::thread(&DatasetImpl<T>::ReleaseMemoryFun, this);
}

template <typename T>
void DatasetImpl<T>::ReleaseMemoryFun() {
572
  VLOG(3) << "DatasetImpl<T>::ReleaseMemory() begin";
J
jiaqi 已提交
573 574 575 576 577 578 579 580 581 582
  if (input_channel_) {
    input_channel_->Clear();
    input_channel_ = nullptr;
  }
  for (size_t i = 0; i < multi_output_channel_.size(); ++i) {
    if (!multi_output_channel_[i]) {
      continue;
    }
    multi_output_channel_[i]->Clear();
    multi_output_channel_[i] = nullptr;
583
  }
J
jiaqi 已提交
584 585 586 587 588 589 590 591 592
  std::vector<paddle::framework::Channel<T>>().swap(multi_output_channel_);
  for (size_t i = 0; i < multi_consume_channel_.size(); ++i) {
    if (!multi_consume_channel_[i]) {
      continue;
    }
    multi_consume_channel_[i]->Clear();
    multi_consume_channel_[i] = nullptr;
  }
  std::vector<paddle::framework::Channel<T>>().swap(multi_consume_channel_);
593 594 595 596
  if (input_pv_channel_) {
    input_pv_channel_->Clear();
    input_pv_channel_ = nullptr;
  }
597 598
  for (auto& pv_output : multi_pv_output_) {
    if (!pv_output) {
599 600
      continue;
    }
601 602
    pv_output->Clear();
    pv_output = nullptr;
603 604
  }
  std::vector<paddle::framework::Channel<PvInstance>>().swap(multi_pv_output_);
605 606
  for (auto& pv_consume : multi_pv_consume_) {
    if (!pv_consume) {
607 608
      continue;
    }
609 610
    pv_consume->Clear();
    pv_consume = nullptr;
611
  }
Y
yaoxuefeng 已提交
612 613 614 615 616 617 618
  if (enable_heterps_) {
    input_records_.clear();
    input_records_.shrink_to_fit();
    std::vector<T>().swap(input_records_);
    VLOG(3) << "release heterps input records records size: "
            << input_records_.size();
  }
619 620
  std::vector<paddle::framework::Channel<PvInstance>>().swap(multi_pv_consume_);

J
jiaqi 已提交
621
  std::vector<std::shared_ptr<paddle::framework::DataFeed>>().swap(readers_);
622 623
  input_records_.clear();
  std::vector<T>().swap(input_records_);
H
hutuxian 已提交
624
  std::vector<T>().swap(slots_shuffle_original_data_);
625
  VLOG(3) << "DatasetImpl<T>::ReleaseMemory() end";
H
hutuxian 已提交
626 627 628 629 630
  VLOG(3) << "total_feasign_num_(" << STAT_GET(STAT_total_feasign_num_in_mem)
          << ") - current_fea_num_(" << total_fea_num_ << ") = ("
          << STAT_GET(STAT_total_feasign_num_in_mem) - total_fea_num_
          << ")";  // For Debug
  STAT_SUB(STAT_total_feasign_num_in_mem, total_fea_num_);
631 632
}

X
xjqbest 已提交
633
// do local shuffle
634 635 636
template <typename T>
void DatasetImpl<T>::LocalShuffle() {
  VLOG(3) << "DatasetImpl<T>::LocalShuffle() begin";
637 638
  platform::Timer timeline;
  timeline.Start();
639

J
jiaqi 已提交
640 641 642
  if (!input_channel_ || input_channel_->Size() == 0) {
    VLOG(3) << "DatasetImpl<T>::LocalShuffle() end, no data to shuffle";
    return;
643
  }
Z
zhaocaibei123 已提交
644
  auto fleet_ptr = framework::FleetWrapper::GetInstance();
J
jiaqi 已提交
645 646 647 648 649 650 651 652 653 654
  input_channel_->Close();
  std::vector<T> data;
  input_channel_->ReadAll(data);
  std::shuffle(data.begin(), data.end(), fleet_ptr->LocalRandomEngine());
  input_channel_->Open();
  input_channel_->Write(std::move(data));
  data.clear();
  data.shrink_to_fit();
  input_channel_->Close();

655 656 657
  timeline.Pause();
  VLOG(3) << "DatasetImpl<T>::LocalShuffle() end, cost time="
          << timeline.ElapsedSec() << " seconds";
658 659
}

660 661 662 663 664 665 666 667 668 669 670 671 672 673 674 675 676 677 678 679
template <typename T>
void DatasetImpl<T>::DumpWalkPath(std::string dump_path, size_t dump_rate) {
  VLOG(3) << "DatasetImpl<T>::DumpWalkPath() begin";
#if defined(PADDLE_WITH_GPU_GRAPH) && defined(PADDLE_WITH_HETERPS)
  std::vector<std::thread> dump_threads;
  if (gpu_graph_mode_) {
    for (int64_t i = 0; i < thread_num_; ++i) {
      dump_threads.push_back(
          std::thread(&paddle::framework::DataFeed::DumpWalkPath,
                      readers_[i].get(),
                      dump_path,
                      dump_rate));
    }
    for (std::thread& t : dump_threads) {
      t.join();
    }
  }
#endif
}

W
wangzhen38 已提交
680 681 682 683 684
// do tdm sample
void MultiSlotDataset::TDMSample(const std::string tree_name,
                                 const std::string tree_path,
                                 const std::vector<uint16_t> tdm_layer_counts,
                                 const uint16_t start_sample_layer,
685 686
                                 const bool with_hierachy,
                                 const uint16_t seed_,
W
wangzhen38 已提交
687 688 689 690 691 692 693
                                 const uint16_t sample_slot) {
#if (defined PADDLE_WITH_DISTRIBUTE) && (defined PADDLE_WITH_PSCORE)
  // init tdm tree
  auto wrapper_ptr = paddle::distributed::IndexWrapper::GetInstance();
  wrapper_ptr->insert_tree_index(tree_name, tree_path);
  auto tree_ptr = wrapper_ptr->get_tree_index(tree_name);
  auto _layer_wise_sample = paddle::distributed::LayerWiseSampler(tree_name);
694 695
  _layer_wise_sample.init_layerwise_conf(
      tdm_layer_counts, start_sample_layer, seed_);
W
wangzhen38 已提交
696 697 698 699 700 701 702 703 704 705 706 707 708 709 710 711 712 713 714 715 716 717 718 719 720 721 722 723 724 725 726 727 728 729

  VLOG(0) << "DatasetImpl<T>::Sample() begin";
  platform::Timer timeline;
  timeline.Start();

  std::vector<std::vector<Record>> data;
  std::vector<std::vector<Record>> sample_results;
  if (!input_channel_ || input_channel_->Size() == 0) {
    for (size_t i = 0; i < multi_output_channel_.size(); ++i) {
      std::vector<Record> tmp_data;
      data.push_back(tmp_data);
      if (!multi_output_channel_[i] || multi_output_channel_[i]->Size() == 0) {
        continue;
      }
      multi_output_channel_[i]->Close();
      multi_output_channel_[i]->ReadAll(data[i]);
    }
  } else {
    input_channel_->Close();
    std::vector<Record> tmp_data;
    data.push_back(tmp_data);
    input_channel_->ReadAll(data[data.size() - 1]);
  }

  VLOG(1) << "finish read src data, data.size = " << data.size()
          << "; details: ";
  auto fleet_ptr = FleetWrapper::GetInstance();
  for (unsigned int i = 0; i < data.size(); i++) {
    VLOG(1) << "data[" << i << "]: size = " << data[i].size();
    std::vector<Record> tmp_results;
    _layer_wise_sample.sample_from_dataset(sample_slot, &data[i], &tmp_results);
    VLOG(1) << "sample_results(" << sample_slot << ") = " << tmp_results.size();
    VLOG(0) << "start to put sample in vector!";
    // sample_results.push_back(tmp_results);
730
    for (auto& tmp_result : tmp_results) {
W
wangzhen38 已提交
731
      std::vector<Record> tmp_vec;
732
      tmp_vec.emplace_back(tmp_result);
W
wangzhen38 已提交
733 734 735 736 737 738
      sample_results.emplace_back(tmp_vec);
    }
    VLOG(0) << "finish to put sample in vector!";
  }

  auto output_channel_num = multi_output_channel_.size();
739
  for (auto& sample_result : sample_results) {
W
wangzhen38 已提交
740 741 742
    auto output_idx = fleet_ptr->LocalRandomEngine()() % output_channel_num;
    multi_output_channel_[output_idx]->Open();
    // vector?
743
    multi_output_channel_[output_idx]->Write(std::move(sample_result));
W
wangzhen38 已提交
744 745 746 747 748 749 750 751 752 753 754 755 756 757
  }

  data.clear();
  sample_results.clear();
  data.shrink_to_fit();
  sample_results.shrink_to_fit();

  timeline.Pause();
  VLOG(0) << "DatasetImpl<T>::Sample() end, cost time=" << timeline.ElapsedSec()
          << " seconds";
#endif
  return;
}

Y
yaoxuefeng 已提交
758 759
void MultiSlotDataset::GlobalShuffle(int thread_num) {
  VLOG(3) << "MultiSlotDataset::GlobalShuffle() begin";
760 761
  platform::Timer timeline;
  timeline.Start();
Z
zhaocaibei123 已提交
762 763 764 765 766
#ifdef PADDLE_WITH_PSCORE
  auto fleet_ptr = distributed::FleetWrapper::GetInstance();
#else
  auto fleet_ptr = framework::FleetWrapper::GetInstance();
#endif
J
jiaqi 已提交
767 768

  if (!input_channel_ || input_channel_->Size() == 0) {
Y
yaoxuefeng 已提交
769
    VLOG(3) << "MultiSlotDataset::GlobalShuffle() end, no data to shuffle";
J
jiaqi 已提交
770 771 772 773 774
    return;
  }

  // local shuffle
  input_channel_->Close();
Y
yaoxuefeng 已提交
775
  std::vector<Record> data;
J
jiaqi 已提交
776 777 778 779 780 781 782 783 784
  input_channel_->ReadAll(data);
  std::shuffle(data.begin(), data.end(), fleet_ptr->LocalRandomEngine());
  input_channel_->Open();
  input_channel_->Write(std::move(data));
  data.clear();
  data.shrink_to_fit();

  input_channel_->Close();
  input_channel_->SetBlockSize(fleet_send_batch_size_);
Y
yaoxuefeng 已提交
785
  VLOG(3) << "MultiSlotDataset::GlobalShuffle() input_channel_ size "
J
jiaqi 已提交
786 787
          << input_channel_->Size();

Y
yaoxuefeng 已提交
788
  auto get_client_id = [this, fleet_ptr](const Record& data) -> size_t {
789
    if (this->merge_by_insid_) {
790 791
      return XXH64(data.ins_id_.data(), data.ins_id_.length(), 0) %
             this->trainer_num_;
792 793 794 795 796
    } else if (this->shuffle_by_uid_) {
      return XXH64(data.uid_.data(), data.uid_.length(), 0) %
             this->trainer_num_;
    } else {
      return fleet_ptr->LocalRandomEngine()() % this->trainer_num_;
797 798 799 800
    }
  };

  auto global_shuffle_func = [this, get_client_id]() {
Z
zhaocaibei123 已提交
801 802 803 804 805 806
#ifdef PADDLE_WITH_PSCORE
    auto fleet_ptr = distributed::FleetWrapper::GetInstance();
#else
    auto fleet_ptr = framework::FleetWrapper::GetInstance();
#endif
    // auto fleet_ptr = framework::FleetWrapper::GetInstance();
Y
yaoxuefeng 已提交
807
    std::vector<Record> data;
J
jiaqi 已提交
808 809 810
    while (this->input_channel_->Read(data)) {
      std::vector<paddle::framework::BinaryArchive> ars(this->trainer_num_);
      for (auto& t : data) {
811
        auto client_id = get_client_id(t);
J
jiaqi 已提交
812 813 814 815 816 817 818
        ars[client_id] << t;
      }
      std::vector<std::future<int32_t>> total_status;
      std::vector<int> send_index(this->trainer_num_);
      for (int i = 0; i < this->trainer_num_; ++i) {
        send_index[i] = i;
      }
819 820
      std::shuffle(
          send_index.begin(), send_index.end(), fleet_ptr->LocalRandomEngine());
821
      for (int index = 0; index < this->trainer_num_; ++index) {
J
jiaqi 已提交
822 823 824 825 826 827 828 829 830 831 832 833 834 835 836
        int i = send_index[index];
        if (ars[i].Length() == 0) {
          continue;
        }
        std::string msg(ars[i].Buffer(), ars[i].Length());
        auto ret = fleet_ptr->SendClientToClientMsg(0, i, msg);
        total_status.push_back(std::move(ret));
      }
      for (auto& t : total_status) {
        t.wait();
      }
      ars.clear();
      ars.shrink_to_fit();
      data.clear();
      data.shrink_to_fit();
837 838 839 840 841 842
      // currently we find bottleneck is server not able to handle large data
      // in time, so we can remove this sleep and set fleet_send_batch_size to
      // 1024, and set server thread to 24.
      if (fleet_send_sleep_seconds_ != 0) {
        sleep(this->fleet_send_sleep_seconds_);
      }
J
jiaqi 已提交
843 844 845
    }
  };

846
  std::vector<std::thread> global_shuffle_threads;
847 848 849 850 851
  if (thread_num == -1) {
    thread_num = thread_num_;
  }
  VLOG(3) << "start global shuffle threads, num = " << thread_num;
  for (int i = 0; i < thread_num; ++i) {
852
    global_shuffle_threads.emplace_back(global_shuffle_func);
853 854 855
  }
  for (std::thread& t : global_shuffle_threads) {
    t.join();
856
  }
J
jiaqi 已提交
857 858 859
  global_shuffle_threads.clear();
  global_shuffle_threads.shrink_to_fit();
  input_channel_->Clear();
860 861 862
  timeline.Pause();
  VLOG(3) << "DatasetImpl<T>::GlobalShuffle() end, cost time="
          << timeline.ElapsedSec() << " seconds";
863 864
}

865
template <typename T>
H
hutuxian 已提交
866 867
void DatasetImpl<T>::DynamicAdjustChannelNum(int channel_num,
                                             bool discard_remaining_ins) {
868 869 870 871 872 873 874 875 876 877
  if (channel_num_ == channel_num) {
    VLOG(3) << "DatasetImpl<T>::DynamicAdjustChannelNum channel_num_="
            << channel_num_ << ", channel_num_=channel_num, no need to adjust";
    return;
  }
  VLOG(3) << "adjust channel num from " << channel_num_ << " to "
          << channel_num;
  channel_num_ = channel_num;
  std::vector<paddle::framework::Channel<T>>* origin_channels = nullptr;
  std::vector<paddle::framework::Channel<T>>* other_channels = nullptr;
878 879 880 881 882
  std::vector<paddle::framework::Channel<PvInstance>>* origin_pv_channels =
      nullptr;
  std::vector<paddle::framework::Channel<PvInstance>>* other_pv_channels =
      nullptr;

883 884 885 886 887
  // find out which channel (output or consume) has data
  int cur_channel = 0;
  uint64_t output_channels_data_size = 0;
  uint64_t consume_channels_data_size = 0;
  CHECK(multi_output_channel_.size() == multi_consume_channel_.size());
888
  for (size_t i = 0; i < multi_output_channel_.size(); ++i) {
889 890 891 892 893 894 895 896 897 898 899 900 901
    output_channels_data_size += multi_output_channel_[i]->Size();
    consume_channels_data_size += multi_consume_channel_[i]->Size();
  }
  if (output_channels_data_size != 0) {
    CHECK(consume_channels_data_size == 0);  // NOLINT
    cur_channel = 0;
  } else {
    CHECK(output_channels_data_size == 0);  // NOLINT
    cur_channel = 1;
  }
  if (cur_channel == 0) {
    origin_channels = &multi_output_channel_;
    other_channels = &multi_consume_channel_;
902 903
    origin_pv_channels = &multi_pv_output_;
    other_pv_channels = &multi_pv_consume_;
904 905 906
  } else {
    origin_channels = &multi_consume_channel_;
    other_channels = &multi_output_channel_;
907 908
    origin_pv_channels = &multi_pv_consume_;
    other_pv_channels = &multi_pv_output_;
909
  }
910 911 912 913
  CHECK(origin_channels != nullptr);     // NOLINT
  CHECK(other_channels != nullptr);      // NOLINT
  CHECK(origin_pv_channels != nullptr);  // NOLINT
  CHECK(other_pv_channels != nullptr);   // NOLINT
914 915 916 917 918

  paddle::framework::Channel<T> total_data_channel =
      paddle::framework::MakeChannel<T>();
  std::vector<paddle::framework::Channel<T>> new_channels;
  std::vector<paddle::framework::Channel<T>> new_other_channels;
919 920 921
  std::vector<paddle::framework::Channel<PvInstance>> new_pv_channels;
  std::vector<paddle::framework::Channel<PvInstance>> new_other_pv_channels;

922
  std::vector<T> local_vec;
923
  for (size_t i = 0; i < origin_channels->size(); ++i) {
924 925 926 927 928 929
    local_vec.clear();
    (*origin_channels)[i]->Close();
    (*origin_channels)[i]->ReadAll(local_vec);
    total_data_channel->Write(std::move(local_vec));
  }
  total_data_channel->Close();
H
hutuxian 已提交
930 931 932 933
  if (static_cast<int>(total_data_channel->Size()) >= channel_num) {
    total_data_channel->SetBlockSize(total_data_channel->Size() / channel_num +
                                     (discard_remaining_ins ? 0 : 1));
  }
H
hutuxian 已提交
934
  if (static_cast<int>(input_channel_->Size()) >= channel_num) {
H
hutuxian 已提交
935 936
    input_channel_->SetBlockSize(input_channel_->Size() / channel_num +
                                 (discard_remaining_ins ? 0 : 1));
H
hutuxian 已提交
937
  }
938 939 940 941 942 943
  if (static_cast<int>(input_pv_channel_->Size()) >= channel_num) {
    input_pv_channel_->SetBlockSize(input_pv_channel_->Size() / channel_num +
                                    (discard_remaining_ins ? 0 : 1));
    VLOG(3) << "now input_pv_channle block size is "
            << input_pv_channel_->BlockSize();
  }
944 945 946 947 948 949 950

  for (int i = 0; i < channel_num; ++i) {
    local_vec.clear();
    total_data_channel->Read(local_vec);
    new_other_channels.push_back(paddle::framework::MakeChannel<T>());
    new_channels.push_back(paddle::framework::MakeChannel<T>());
    new_channels[i]->Write(std::move(local_vec));
951 952 953
    new_other_pv_channels.push_back(
        paddle::framework::MakeChannel<PvInstance>());
    new_pv_channels.push_back(paddle::framework::MakeChannel<PvInstance>());
954 955 956 957 958 959 960 961
  }

  total_data_channel->Clear();
  origin_channels->clear();
  other_channels->clear();
  *origin_channels = new_channels;
  *other_channels = new_other_channels;

962 963 964 965 966
  origin_pv_channels->clear();
  other_pv_channels->clear();
  *origin_pv_channels = new_pv_channels;
  *other_pv_channels = new_other_pv_channels;

967 968 969 970
  new_channels.clear();
  new_other_channels.clear();
  std::vector<paddle::framework::Channel<T>>().swap(new_channels);
  std::vector<paddle::framework::Channel<T>>().swap(new_other_channels);
971 972 973 974 975 976 977

  new_pv_channels.clear();
  new_other_pv_channels.clear();
  std::vector<paddle::framework::Channel<PvInstance>>().swap(new_pv_channels);
  std::vector<paddle::framework::Channel<PvInstance>>().swap(
      new_other_pv_channels);

978 979 980 981 982 983 984 985 986 987 988 989 990 991 992 993 994 995 996 997 998 999 1000 1001
  local_vec.clear();
  std::vector<T>().swap(local_vec);
  VLOG(3) << "adjust channel num done";
}

template <typename T>
void DatasetImpl<T>::DynamicAdjustReadersNum(int thread_num) {
  if (thread_num_ == thread_num) {
    VLOG(3) << "DatasetImpl<T>::DynamicAdjustReadersNum thread_num_="
            << thread_num_ << ", thread_num_=thread_num, no need to adjust";
    return;
  }
  VLOG(3) << "adjust readers num from " << thread_num_ << " to " << thread_num;
  thread_num_ = thread_num;
  std::vector<std::shared_ptr<paddle::framework::DataFeed>>().swap(readers_);
  CreateReaders();
  VLOG(3) << "adjust readers num done";
}

template <typename T>
void DatasetImpl<T>::SetFleetSendSleepSeconds(int seconds) {
  fleet_send_sleep_seconds_ = seconds;
}

1002 1003
template <typename T>
void DatasetImpl<T>::CreateReaders() {
1004
  VLOG(3) << "Calling CreateReaders()";
J
jiaqi 已提交
1005 1006 1007 1008 1009 1010
  VLOG(3) << "thread num in Dataset: " << thread_num_;
  VLOG(3) << "Filelist size in Dataset: " << filelist_.size();
  VLOG(3) << "channel num in Dataset: " << channel_num_;
  CHECK(thread_num_ > 0) << "thread num should > 0";
  CHECK(channel_num_ > 0) << "channel num should > 0";
  CHECK(channel_num_ <= thread_num_) << "channel num should <= thread num";
1011
  VLOG(3) << "readers size: " << readers_.size();
1012
  if (!readers_.empty()) {
J
jiaqi 已提交
1013 1014
    VLOG(3) << "readers_.size() = " << readers_.size()
            << ", will not create again";
1015 1016
    return;
  }
1017
  VLOG(3) << "data feed class name: " << data_feed_desc_.name();
J
jiaqi 已提交
1018
  int channel_idx = 0;
1019
  for (int i = 0; i < thread_num_; ++i) {
1020
    readers_.push_back(DataFeedFactory::CreateDataFeed(data_feed_desc_.name()));
J
jiaqi 已提交
1021 1022 1023 1024 1025
    readers_[i]->Init(data_feed_desc_);
    readers_[i]->SetThreadId(i);
    readers_[i]->SetThreadNum(thread_num_);
    readers_[i]->SetFileListMutex(&mutex_for_pick_file_);
    readers_[i]->SetFileListIndex(&file_idx_);
H
hutuxian 已提交
1026 1027
    readers_[i]->SetFeaNumMutex(&mutex_for_fea_num_);
    readers_[i]->SetFeaNum(&total_fea_num_);
J
jiaqi 已提交
1028
    readers_[i]->SetFileList(filelist_);
1029
    readers_[i]->SetParseInsId(parse_ins_id_);
1030
    readers_[i]->SetParseUid(parse_uid_);
1031
    readers_[i]->SetParseContent(parse_content_);
1032 1033 1034 1035 1036 1037
    readers_[i]->SetParseLogKey(parse_logkey_);
    readers_[i]->SetEnablePvMerge(enable_pv_merge_);
    // Notice: it is only valid for untest of test_paddlebox_datafeed.
    // In fact, it does not affect the train process when paddle is
    // complied with Box_Ps.
    readers_[i]->SetCurrentPhase(current_phase_);
J
jiaqi 已提交
1038 1039 1040
    if (input_channel_ != nullptr) {
      readers_[i]->SetInputChannel(input_channel_.get());
    }
1041 1042 1043
    if (input_pv_channel_ != nullptr) {
      readers_[i]->SetInputPvChannel(input_pv_channel_.get());
    }
1044 1045
    if (cur_channel_ == 0 &&
        static_cast<size_t>(channel_idx) < multi_output_channel_.size()) {
J
jiaqi 已提交
1046 1047
      readers_[i]->SetOutputChannel(multi_output_channel_[channel_idx].get());
      readers_[i]->SetConsumeChannel(multi_consume_channel_[channel_idx].get());
1048 1049
      readers_[i]->SetOutputPvChannel(multi_pv_output_[channel_idx].get());
      readers_[i]->SetConsumePvChannel(multi_pv_consume_[channel_idx].get());
1050 1051
    } else if (static_cast<size_t>(channel_idx) <
               multi_output_channel_.size()) {
J
jiaqi 已提交
1052 1053
      readers_[i]->SetOutputChannel(multi_consume_channel_[channel_idx].get());
      readers_[i]->SetConsumeChannel(multi_output_channel_[channel_idx].get());
1054 1055
      readers_[i]->SetOutputPvChannel(multi_pv_consume_[channel_idx].get());
      readers_[i]->SetConsumePvChannel(multi_pv_output_[channel_idx].get());
J
jiaqi 已提交
1056 1057 1058 1059 1060
    }
    ++channel_idx;
    if (channel_idx >= channel_num_) {
      channel_idx = 0;
    }
1061
  }
J
jiaqi 已提交
1062
  VLOG(3) << "readers size: " << readers_.size();
1063 1064
}

1065 1066 1067
template <typename T>
void DatasetImpl<T>::DestroyReaders() {
  VLOG(3) << "Calling DestroyReaders()";
1068
  VLOG(3) << "readers size1: " << readers_.size();
1069
  std::vector<std::shared_ptr<paddle::framework::DataFeed>>().swap(readers_);
1070
  VLOG(3) << "readers size: " << readers_.size();
J
jiaqi 已提交
1071 1072
  file_idx_ = 0;
  cur_channel_ = 1 - cur_channel_;
1073 1074
}

1075 1076 1077 1078 1079 1080 1081 1082 1083 1084 1085 1086 1087 1088 1089 1090 1091 1092 1093 1094 1095 1096 1097
template <typename T>
void DatasetImpl<T>::SetPreLoadThreadNum(int thread_num) {
  preload_thread_num_ = thread_num;
}

template <typename T>
void DatasetImpl<T>::CreatePreLoadReaders() {
  VLOG(3) << "Begin CreatePreLoadReaders";
  if (preload_thread_num_ == 0) {
    preload_thread_num_ = thread_num_;
  }
  CHECK(preload_thread_num_ > 0) << "thread num should > 0";
  CHECK(input_channel_ != nullptr);
  preload_readers_.clear();
  for (int i = 0; i < preload_thread_num_; ++i) {
    preload_readers_.push_back(
        DataFeedFactory::CreateDataFeed(data_feed_desc_.name()));
    preload_readers_[i]->Init(data_feed_desc_);
    preload_readers_[i]->SetThreadId(i);
    preload_readers_[i]->SetThreadNum(preload_thread_num_);
    preload_readers_[i]->SetFileListMutex(&mutex_for_pick_file_);
    preload_readers_[i]->SetFileListIndex(&file_idx_);
    preload_readers_[i]->SetFileList(filelist_);
H
hutuxian 已提交
1098 1099
    preload_readers_[i]->SetFeaNumMutex(&mutex_for_fea_num_);
    preload_readers_[i]->SetFeaNum(&total_fea_num_);
1100
    preload_readers_[i]->SetParseInsId(parse_ins_id_);
1101
    preload_readers_[i]->SetParseUid(parse_uid_);
1102
    preload_readers_[i]->SetParseContent(parse_content_);
1103 1104
    preload_readers_[i]->SetParseLogKey(parse_logkey_);
    preload_readers_[i]->SetEnablePvMerge(enable_pv_merge_);
1105 1106 1107
    preload_readers_[i]->SetInputChannel(input_channel_.get());
    preload_readers_[i]->SetOutputChannel(nullptr);
    preload_readers_[i]->SetConsumeChannel(nullptr);
1108 1109
    preload_readers_[i]->SetOutputPvChannel(nullptr);
    preload_readers_[i]->SetConsumePvChannel(nullptr);
1110 1111 1112 1113 1114 1115 1116 1117 1118 1119 1120 1121 1122 1123
  }
  VLOG(3) << "End CreatePreLoadReaders";
}

template <typename T>
void DatasetImpl<T>::DestroyPreLoadReaders() {
  VLOG(3) << "Begin DestroyPreLoadReaders";
  preload_readers_.clear();
  std::vector<std::shared_ptr<paddle::framework::DataFeed>>().swap(
      preload_readers_);
  file_idx_ = 0;
  VLOG(3) << "End DestroyPreLoadReaders";
}

1124 1125
template <typename T>
int64_t DatasetImpl<T>::GetMemoryDataSize() {
L
lxsbupt 已提交
1126 1127 1128 1129 1130 1131 1132 1133 1134 1135 1136 1137 1138 1139 1140 1141 1142 1143 1144 1145 1146 1147 1148 1149
  if (gpu_graph_mode_) {
    int64_t total_path_num = 0;
    for (int i = 0; i < thread_num_; i++) {
      total_path_num += readers_[i]->GetGraphPathNum();
    }
    return total_path_num;
  } else {
    return input_channel_->Size();
  }
}

template <typename T>
bool DatasetImpl<T>::GetEpochFinish() {
#if defined(PADDLE_WITH_GPU_GRAPH) && defined(PADDLE_WITH_HETERPS)
  bool is_epoch_finish = true;
  if (gpu_graph_mode_) {
    for (int i = 0; i < thread_num_; i++) {
      is_epoch_finish = is_epoch_finish && readers_[i]->get_epoch_finish();
    }
  }
  return is_epoch_finish;
#else
  return false;
#endif
1150 1151
}

1152 1153 1154 1155 1156 1157 1158 1159 1160 1161
template <typename T>
int64_t DatasetImpl<T>::GetPvDataSize() {
  if (enable_pv_merge_) {
    return input_pv_channel_->Size();
  } else {
    VLOG(0) << "It does not merge pv..";
    return 0;
  }
}

1162 1163 1164
template <typename T>
int64_t DatasetImpl<T>::GetShuffleDataSize() {
  int64_t sum = 0;
J
jiaqi 已提交
1165 1166
  for (size_t i = 0; i < multi_output_channel_.size(); ++i) {
    sum += multi_output_channel_[i]->Size() + multi_consume_channel_[i]->Size();
1167 1168 1169 1170
  }
  return sum;
}

1171 1172
int MultiSlotDataset::ReceiveFromClient(int msg_type,
                                        int client_id,
Y
yaoxuefeng 已提交
1173
                                        const std::string& msg) {
D
dongdaxiang 已提交
1174
#ifdef _LINUX
1175
  VLOG(3) << "ReceiveFromClient msg_type=" << msg_type
1176
          << ", client_id=" << client_id << ", msg length=" << msg.length();
J
jiaqi 已提交
1177 1178 1179 1180 1181 1182 1183 1184
  if (msg.length() == 0) {
    return 0;
  }
  paddle::framework::BinaryArchive ar;
  ar.SetReadBuffer(const_cast<char*>(msg.c_str()), msg.length(), nullptr);
  if (ar.Cursor() == ar.Finish()) {
    return 0;
  }
Y
yaoxuefeng 已提交
1185
  std::vector<Record> data;
J
jiaqi 已提交
1186
  while (ar.Cursor() < ar.Finish()) {
Y
yaoxuefeng 已提交
1187
    data.push_back(ar.Get<Record>());
J
jiaqi 已提交
1188 1189 1190
  }
  CHECK(ar.Cursor() == ar.Finish());

Z
zhaocaibei123 已提交
1191
  auto fleet_ptr = framework::FleetWrapper::GetInstance();
1192 1193 1194 1195 1196 1197 1198 1199 1200 1201
  // not use random because it doesn't perform well here.
  // to make sure each channel get data equally, we just put data to
  // channel one by one.
  // int64_t index = fleet_ptr->LocalRandomEngine()() % channel_num_;
  int64_t index = 0;
  {
    std::unique_lock<std::mutex> lk(global_index_mutex_);
    index = global_index_++;
  }
  index = index % channel_num_;
1202
  VLOG(3) << "ramdom index=" << index;
J
jiaqi 已提交
1203 1204 1205 1206
  multi_output_channel_[index]->Write(std::move(data));

  data.clear();
  data.shrink_to_fit();
D
dongdaxiang 已提交
1207
#endif
1208 1209 1210
  return 0;
}

1211
// explicit instantiation
J
jiaqi 已提交
1212
template class DatasetImpl<Record>;
1213

Y
yaoxuefeng 已提交
1214 1215 1216 1217 1218 1219 1220 1221 1222 1223 1224 1225 1226 1227
void MultiSlotDataset::DynamicAdjustReadersNum(int thread_num) {
  if (thread_num_ == thread_num) {
    VLOG(3) << "DatasetImpl<T>::DynamicAdjustReadersNum thread_num_="
            << thread_num_ << ", thread_num_=thread_num, no need to adjust";
    return;
  }
  VLOG(3) << "adjust readers num from " << thread_num_ << " to " << thread_num;
  thread_num_ = thread_num;
  std::vector<std::shared_ptr<paddle::framework::DataFeed>>().swap(readers_);
  CreateReaders();
  VLOG(3) << "adjust readers num done";
  PrepareTrain();
}

1228 1229 1230
void MultiSlotDataset::PostprocessInstance() {
  // divide pv instance, and merge to input_channel_
  if (enable_pv_merge_) {
Z
zhaocaibei123 已提交
1231
    auto fleet_ptr = framework::FleetWrapper::GetInstance();
1232 1233
    std::shuffle(input_records_.begin(),
                 input_records_.end(),
1234
                 fleet_ptr->LocalRandomEngine());
1235 1236
    input_channel_->Open();
    input_channel_->Write(std::move(input_records_));
1237 1238
    for (auto& pv_consume : multi_pv_consume_) {
      pv_consume->Clear();
1239 1240 1241 1242 1243 1244
    }
    input_channel_->Close();
    input_records_.clear();
    input_records_.shrink_to_fit();
  } else {
    input_channel_->Open();
1245
    for (auto& consume_channel : multi_consume_channel_) {
1246
      std::vector<Record> ins_data;
1247 1248
      consume_channel->Close();
      consume_channel->ReadAll(ins_data);
1249 1250 1251
      input_channel_->Write(std::move(ins_data));
      ins_data.clear();
      ins_data.shrink_to_fit();
1252
      consume_channel->Clear();
1253 1254
    }
    input_channel_->Close();
1255
    this->LocalShuffle();
1256 1257 1258 1259 1260 1261 1262 1263 1264 1265 1266 1267 1268 1269
  }
}

void MultiSlotDataset::SetCurrentPhase(int current_phase) {
  current_phase_ = current_phase;
}

void MultiSlotDataset::PreprocessInstance() {
  if (!input_channel_ || input_channel_->Size() == 0) {
    return;
  }
  if (!enable_pv_merge_) {  // means to use Record
    this->LocalShuffle();
  } else {  // means to use Pv
Z
zhaocaibei123 已提交
1270
    auto fleet_ptr = framework::FleetWrapper::GetInstance();
1271 1272 1273 1274 1275 1276 1277 1278 1279 1280
    input_channel_->Close();
    std::vector<PvInstance> pv_data;
    input_channel_->ReadAll(input_records_);
    int all_records_num = input_records_.size();
    std::vector<Record*> all_records;
    all_records.reserve(all_records_num);
    for (int index = 0; index < all_records_num; ++index) {
      all_records.push_back(&input_records_[index]);
    }

1281 1282
    std::sort(all_records.data(),
              all_records.data() + all_records_num,
1283 1284 1285 1286 1287 1288 1289 1290 1291 1292 1293 1294 1295 1296 1297 1298 1299 1300 1301 1302 1303 1304 1305 1306 1307
              [](const Record* lhs, const Record* rhs) {
                return lhs->search_id < rhs->search_id;
              });
    if (merge_by_sid_) {
      uint64_t last_search_id = 0;
      for (int i = 0; i < all_records_num; ++i) {
        Record* ins = all_records[i];
        if (i == 0 || last_search_id != ins->search_id) {
          PvInstance pv_instance = make_pv_instance();
          pv_instance->merge_instance(ins);
          pv_data.push_back(pv_instance);
          last_search_id = ins->search_id;
          continue;
        }
        pv_data.back()->merge_instance(ins);
      }
    } else {
      for (int i = 0; i < all_records_num; ++i) {
        Record* ins = all_records[i];
        PvInstance pv_instance = make_pv_instance();
        pv_instance->merge_instance(ins);
        pv_data.push_back(pv_instance);
      }
    }

1308 1309
    std::shuffle(
        pv_data.begin(), pv_data.end(), fleet_ptr->LocalRandomEngine());
1310 1311 1312 1313 1314 1315 1316 1317 1318
    input_pv_channel_->Open();
    input_pv_channel_->Write(std::move(pv_data));

    pv_data.clear();
    pv_data.shrink_to_fit();
    input_pv_channel_->Close();
  }
}

1319 1320
void MultiSlotDataset::GenerateLocalTablesUnlock(int table_id,
                                                 int feadim,
1321 1322 1323 1324 1325 1326 1327 1328 1329 1330
                                                 int read_thread_num,
                                                 int consume_thread_num,
                                                 int shard_num) {
  VLOG(3) << "MultiSlotDataset::GenerateUniqueFeasign begin";
  if (!gen_uni_feasigns_) {
    VLOG(3) << "generate_unique_feasign_=false, will not GenerateUniqueFeasign";
    return;
  }

  CHECK(multi_output_channel_.size() != 0);  // NOLINT
Z
zhaocaibei123 已提交
1331
  auto fleet_ptr_ = framework::FleetWrapper::GetInstance();
1332 1333 1334 1335 1336 1337 1338 1339 1340 1341
  std::vector<std::unordered_map<uint64_t, std::vector<float>>>&
      local_map_tables = fleet_ptr_->GetLocalTable();
  local_map_tables.resize(shard_num);
  // read thread
  int channel_num = multi_output_channel_.size();
  if (read_thread_num < channel_num) {
    read_thread_num = channel_num;
  }
  std::vector<std::thread> threads(read_thread_num);
  consume_task_pool_.resize(consume_thread_num);
1342 1343
  for (auto& consume_task : consume_task_pool_) {
    consume_task.reset(new ::ThreadPool(1));
1344
  }
1345 1346
  auto consume_func = [&local_map_tables](int shard_id,
                                          int feadim,
1347 1348 1349 1350 1351 1352 1353 1354
                                          std::vector<uint64_t>& keys) {
    for (auto k : keys) {
      if (local_map_tables[shard_id].find(k) ==
          local_map_tables[shard_id].end()) {
        local_map_tables[shard_id][k] = std::vector<float>(feadim, 0);
      }
    }
  };
1355 1356 1357 1358 1359 1360 1361
  auto gen_func =
      [this, &shard_num, &feadim, &local_map_tables, &consume_func](int i) {
        std::vector<Record> vec_data;
        std::vector<std::vector<uint64_t>> task_keys(shard_num);
        std::vector<std::future<void>> task_futures;
        this->multi_output_channel_[i]->Close();
        this->multi_output_channel_[i]->ReadAll(vec_data);
1362 1363
        for (auto& item : vec_data) {
          for (auto& feature : item.uint64_feasigns_) {
1364 1365 1366 1367
            int shard = feature.sign().uint64_feasign_ % shard_num;
            task_keys[shard].push_back(feature.sign().uint64_feasign_);
          }
        }
1368

1369 1370 1371 1372
        for (int shard_id = 0; shard_id < shard_num; shard_id++) {
          task_futures.emplace_back(consume_task_pool_[shard_id]->enqueue(
              consume_func, shard_id, feadim, task_keys[shard_id]));
        }
1373

1374 1375 1376 1377 1378 1379 1380 1381 1382 1383 1384 1385 1386 1387
        multi_output_channel_[i]->Open();
        multi_output_channel_[i]->Write(std::move(vec_data));
        vec_data.clear();
        vec_data.shrink_to_fit();
        for (auto& tk : task_keys) {
          tk.clear();
          std::vector<uint64_t>().swap(tk);
        }
        task_keys.clear();
        std::vector<std::vector<uint64_t>>().swap(task_keys);
        for (auto& tf : task_futures) {
          tf.wait();
        }
      };
1388 1389 1390 1391 1392 1393
  for (size_t i = 0; i < threads.size(); i++) {
    threads[i] = std::thread(gen_func, i);
  }
  for (std::thread& t : threads) {
    t.join();
  }
1394 1395
  for (auto& consume_task : consume_task_pool_) {
    consume_task.reset();
1396 1397 1398 1399
  }
  consume_task_pool_.clear();
  fleet_ptr_->PullSparseToLocal(table_id, feadim);
}
1400

1401 1402 1403 1404 1405 1406 1407 1408
void MultiSlotDataset::MergeByInsId() {
  VLOG(3) << "MultiSlotDataset::MergeByInsId begin";
  if (!merge_by_insid_) {
    VLOG(3) << "merge_by_insid=false, will not MergeByInsId";
    return;
  }
  auto multi_slot_desc = data_feed_desc_.multi_slot_desc();
  std::vector<std::string> use_slots;
1409
  std::vector<bool> use_slots_is_dense;
1410
  for (int i = 0; i < multi_slot_desc.slots_size(); ++i) {
1411 1412 1413
    const auto& slot = multi_slot_desc.slots(i);
    if (slot.is_used()) {
      use_slots.push_back(slot.name());
1414
      use_slots_is_dense.push_back(slot.is_dense());
1415 1416 1417 1418 1419
    }
  }
  CHECK(multi_output_channel_.size() != 0);  // NOLINT
  auto channel_data = paddle::framework::MakeChannel<Record>();
  VLOG(3) << "multi_output_channel_.size() " << multi_output_channel_.size();
1420
  for (auto& item : multi_output_channel_) {
1421
    std::vector<Record> vec_data;
1422 1423
    item->Close();
    item->ReadAll(vec_data);
1424 1425 1426
    channel_data->Write(std::move(vec_data));
    vec_data.clear();
    vec_data.shrink_to_fit();
1427
    item->Clear();
1428 1429 1430 1431 1432 1433 1434 1435 1436 1437 1438
  }
  channel_data->Close();
  std::vector<Record> recs;
  recs.reserve(channel_data->Size());
  channel_data->ReadAll(recs);
  channel_data->Clear();
  std::sort(recs.begin(), recs.end(), [](const Record& a, const Record& b) {
    return a.ins_id_ < b.ins_id_;
  });

  std::vector<Record> results;
1439 1440 1441 1442 1443
  uint64_t drop_ins_num = 0;
  std::unordered_set<uint16_t> all_int64;
  std::unordered_set<uint16_t> all_float;
  std::unordered_set<uint16_t> local_uint64;
  std::unordered_set<uint16_t> local_float;
1444 1445 1446 1447 1448
  std::unordered_map<uint16_t, std::vector<FeatureItem>> all_dense_uint64;
  std::unordered_map<uint16_t, std::vector<FeatureItem>> all_dense_float;
  std::unordered_map<uint16_t, std::vector<FeatureItem>> local_dense_uint64;
  std::unordered_map<uint16_t, std::vector<FeatureItem>> local_dense_float;
  std::unordered_map<uint16_t, bool> dense_empty;
1449

1450 1451 1452 1453 1454 1455
  VLOG(3) << "recs.size() " << recs.size();
  for (size_t i = 0; i < recs.size();) {
    size_t j = i + 1;
    while (j < recs.size() && recs[j].ins_id_ == recs[i].ins_id_) {
      j++;
    }
1456 1457 1458 1459
    if (merge_size_ > 0 && j - i != merge_size_) {
      drop_ins_num += j - i;
      LOG(WARNING) << "drop ins " << recs[i].ins_id_ << " size=" << j - i
                   << ", because merge_size=" << merge_size_;
1460 1461 1462 1463
      i = j;
      continue;
    }

1464 1465
    all_int64.clear();
    all_float.clear();
1466 1467
    all_dense_uint64.clear();
    all_dense_float.clear();
1468 1469 1470 1471 1472 1473
    bool has_conflict_slot = false;
    uint16_t conflict_slot = 0;

    Record rec;
    rec.ins_id_ = recs[i].ins_id_;
    rec.content_ = recs[i].content_;
1474

1475 1476 1477 1478 1479 1480 1481 1482 1483 1484 1485 1486 1487 1488 1489 1490 1491 1492 1493 1494 1495 1496 1497 1498 1499 1500 1501 1502 1503 1504 1505 1506 1507 1508 1509 1510 1511 1512 1513
    for (size_t k = i; k < j; k++) {
      dense_empty.clear();
      local_dense_uint64.clear();
      local_dense_float.clear();
      for (auto& feature : recs[k].uint64_feasigns_) {
        uint16_t slot = feature.slot();
        if (!use_slots_is_dense[slot]) {
          continue;
        }
        local_dense_uint64[slot].push_back(feature);
        if (feature.sign().uint64_feasign_ != 0) {
          dense_empty[slot] = false;
        } else if (dense_empty.find(slot) == dense_empty.end() &&
                   all_dense_uint64.find(slot) == all_dense_uint64.end()) {
          dense_empty[slot] = true;
        }
      }
      for (auto& feature : recs[k].float_feasigns_) {
        uint16_t slot = feature.slot();
        if (!use_slots_is_dense[slot]) {
          continue;
        }
        local_dense_float[slot].push_back(feature);
        if (fabs(feature.sign().float_feasign_) >= 1e-6) {
          dense_empty[slot] = false;
        } else if (dense_empty.find(slot) == dense_empty.end() &&
                   all_dense_float.find(slot) == all_dense_float.end()) {
          dense_empty[slot] = true;
        }
      }
      for (auto& p : dense_empty) {
        if (local_dense_uint64.find(p.first) != local_dense_uint64.end()) {
          all_dense_uint64[p.first] = std::move(local_dense_uint64[p.first]);
        } else if (local_dense_float.find(p.first) != local_dense_float.end()) {
          all_dense_float[p.first] = std::move(local_dense_float[p.first]);
        }
      }
    }
    for (auto& f : all_dense_uint64) {
1514 1515
      rec.uint64_feasigns_.insert(
          rec.uint64_feasigns_.end(), f.second.begin(), f.second.end());
1516 1517
    }
    for (auto& f : all_dense_float) {
1518 1519
      rec.float_feasigns_.insert(
          rec.float_feasigns_.end(), f.second.begin(), f.second.end());
1520 1521
    }

1522 1523 1524
    for (size_t k = i; k < j; k++) {
      local_uint64.clear();
      local_float.clear();
1525
      for (auto& feature : recs[k].uint64_feasigns_) {
1526
        uint16_t slot = feature.slot();
1527 1528 1529
        if (use_slots_is_dense[slot]) {
          continue;
        } else if (all_int64.find(slot) != all_int64.end()) {
1530 1531 1532
          has_conflict_slot = true;
          conflict_slot = slot;
          break;
1533
        }
1534 1535 1536 1537 1538
        local_uint64.insert(slot);
        rec.uint64_feasigns_.push_back(std::move(feature));
      }
      if (has_conflict_slot) {
        break;
1539
      }
1540 1541
      all_int64.insert(local_uint64.begin(), local_uint64.end());

1542
      for (auto& feature : recs[k].float_feasigns_) {
1543
        uint16_t slot = feature.slot();
1544 1545 1546
        if (use_slots_is_dense[slot]) {
          continue;
        } else if (all_float.find(slot) != all_float.end()) {
1547 1548 1549
          has_conflict_slot = true;
          conflict_slot = slot;
          break;
1550
        }
1551 1552 1553 1554 1555
        local_float.insert(slot);
        rec.float_feasigns_.push_back(std::move(feature));
      }
      if (has_conflict_slot) {
        break;
1556
      }
1557
      all_float.insert(local_float.begin(), local_float.end());
1558 1559
    }

1560 1561 1562 1563
    if (has_conflict_slot) {
      LOG(WARNING) << "drop ins " << recs[i].ins_id_ << " size=" << j - i
                   << ", because conflict_slot=" << use_slots[conflict_slot];
      drop_ins_num += j - i;
1564
    } else {
1565
      results.push_back(std::move(rec));
1566
    }
1567
    i = j;
1568
  }
1569
  std::vector<Record>().swap(recs);
1570
  VLOG(3) << "results size " << results.size();
1571
  LOG(WARNING) << "total drop ins num: " << drop_ins_num;
1572 1573
  results.shrink_to_fit();

Z
zhaocaibei123 已提交
1574
  auto fleet_ptr = framework::FleetWrapper::GetInstance();
1575 1576 1577 1578 1579 1580 1581 1582 1583
  std::shuffle(results.begin(), results.end(), fleet_ptr->LocalRandomEngine());
  channel_data->Open();
  channel_data->Write(std::move(results));
  channel_data->Close();
  results.clear();
  results.shrink_to_fit();
  VLOG(3) << "channel data size " << channel_data->Size();
  channel_data->SetBlockSize(channel_data->Size() / channel_num_ + 1);
  VLOG(3) << "channel data block size " << channel_data->BlockSize();
1584
  for (auto& item : multi_output_channel_) {
1585 1586
    std::vector<Record> vec_data;
    channel_data->Read(vec_data);
1587 1588
    item->Open();
    item->Write(std::move(vec_data));
1589 1590 1591 1592 1593 1594 1595 1596
    vec_data.clear();
    vec_data.shrink_to_fit();
  }
  CHECK(channel_data->Size() == 0);  // NOLINT
  channel_data->Clear();
  VLOG(3) << "MultiSlotDataset::MergeByInsId end";
}

1597 1598 1599
void MultiSlotDataset::GetRandomData(
    const std::unordered_set<uint16_t>& slots_to_replace,
    std::vector<Record>* result) {
1600 1601 1602 1603
  int debug_erase_cnt = 0;
  int debug_push_cnt = 0;
  auto multi_slot_desc = data_feed_desc_.multi_slot_desc();
  slots_shuffle_rclist_.ReInit();
1604 1605
  const auto& slots_shuffle_original_data = GetSlotsOriginalData();
  for (const auto& rec : slots_shuffle_original_data) {
1606 1607 1608 1609 1610 1611 1612 1613 1614 1615 1616 1617 1618
    RecordCandidate rand_rec;
    Record new_rec = rec;
    slots_shuffle_rclist_.AddAndGet(rec, &rand_rec);
    for (auto it = new_rec.uint64_feasigns_.begin();
         it != new_rec.uint64_feasigns_.end();) {
      if (slots_to_replace.find(it->slot()) != slots_to_replace.end()) {
        it = new_rec.uint64_feasigns_.erase(it);
        debug_erase_cnt += 1;
      } else {
        ++it;
      }
    }
    for (auto slot : slots_to_replace) {
1619
      auto range = rand_rec.feas_.equal_range(slot);
1620
      for (auto it = range.first; it != range.second; ++it) {
1621
        new_rec.uint64_feasigns_.emplace_back(it->second, it->first);
1622 1623 1624 1625 1626 1627 1628 1629 1630
        debug_push_cnt += 1;
      }
    }
    result->push_back(std::move(new_rec));
  }
  VLOG(2) << "erase feasign num: " << debug_erase_cnt
          << " repush feasign num: " << debug_push_cnt;
}

1631 1632 1633
void MultiSlotDataset::PreprocessChannel(
    const std::set<std::string>& slots_to_replace,
    std::unordered_set<uint16_t>& index_slots) {  // NOLINT
1634 1635
  int out_channel_size = 0;
  if (cur_channel_ == 0) {
1636 1637
    for (auto& item : multi_output_channel_) {
      out_channel_size += item->Size();
1638 1639
    }
  } else {
1640 1641
    for (auto& item : multi_consume_channel_) {
      out_channel_size += item->Size();
1642 1643 1644 1645 1646
    }
  }
  VLOG(2) << "DatasetImpl<T>::SlotsShuffle() begin with input channel size: "
          << input_channel_->Size()
          << " output channel size: " << out_channel_size;
1647

1648
  if ((!input_channel_ || input_channel_->Size() == 0) &&
1649
      slots_shuffle_original_data_.empty() && out_channel_size == 0) {
1650 1651 1652
    VLOG(3) << "DatasetImpl<T>::SlotsShuffle() end, no data to slots shuffle";
    return;
  }
1653

1654
  auto multi_slot_desc = data_feed_desc_.multi_slot_desc();
1655
  for (int i = 0; i < multi_slot_desc.slots_size(); ++i) {
1656 1657 1658 1659 1660
    std::string cur_slot = multi_slot_desc.slots(i).name();
    if (slots_to_replace.find(cur_slot) != slots_to_replace.end()) {
      index_slots.insert(i);
    }
  }
1661
  if (slots_shuffle_original_data_.empty()) {
1662 1663 1664 1665 1666 1667 1668 1669 1670
    // before first slots shuffle, instances could be in
    // input_channel, oupput_channel or consume_channel
    if (input_channel_ && input_channel_->Size() != 0) {
      slots_shuffle_original_data_.reserve(input_channel_->Size());
      input_channel_->Close();
      input_channel_->ReadAll(slots_shuffle_original_data_);
    } else {
      CHECK(out_channel_size > 0);  // NOLINT
      if (cur_channel_ == 0) {
1671
        for (auto& item : multi_output_channel_) {
1672
          std::vector<Record> vec_data;
1673 1674
          item->Close();
          item->ReadAll(vec_data);
1675 1676 1677 1678 1679 1680 1681 1682
          slots_shuffle_original_data_.reserve(
              slots_shuffle_original_data_.size() + vec_data.size());
          slots_shuffle_original_data_.insert(
              slots_shuffle_original_data_.end(),
              std::make_move_iterator(vec_data.begin()),
              std::make_move_iterator(vec_data.end()));
          vec_data.clear();
          vec_data.shrink_to_fit();
1683
          item->Clear();
1684 1685
        }
      } else {
1686
        for (auto& item : multi_consume_channel_) {
1687
          std::vector<Record> vec_data;
1688 1689
          item->Close();
          item->ReadAll(vec_data);
1690 1691 1692 1693 1694 1695 1696 1697
          slots_shuffle_original_data_.reserve(
              slots_shuffle_original_data_.size() + vec_data.size());
          slots_shuffle_original_data_.insert(
              slots_shuffle_original_data_.end(),
              std::make_move_iterator(vec_data.begin()),
              std::make_move_iterator(vec_data.end()));
          vec_data.clear();
          vec_data.shrink_to_fit();
1698
          item->Clear();
1699 1700 1701 1702 1703 1704 1705
        }
      }
    }
  } else {
    // if already have original data for slots shuffle, clear channel
    input_channel_->Clear();
    if (cur_channel_ == 0) {
1706 1707
      for (auto& item : multi_output_channel_) {
        if (!item) {
1708 1709
          continue;
        }
1710
        item->Clear();
1711 1712
      }
    } else {
1713 1714
      for (auto& item : multi_consume_channel_) {
        if (!item) {
1715 1716
          continue;
        }
1717
        item->Clear();
1718 1719 1720 1721 1722
      }
    }
  }
  int end_size = 0;
  if (cur_channel_ == 0) {
1723 1724
    for (auto& item : multi_output_channel_) {
      if (!item) {
1725 1726
        continue;
      }
1727
      end_size += item->Size();
1728 1729
    }
  } else {
1730 1731
    for (auto& item : multi_consume_channel_) {
      if (!item) {
1732 1733
        continue;
      }
1734
      end_size += item->Size();
1735 1736 1737 1738
    }
  }
  CHECK(input_channel_->Size() == 0)
      << "input channel should be empty before slots shuffle";
1739 1740 1741 1742 1743
}

// slots shuffle to input_channel_ with needed-shuffle slots
void MultiSlotDataset::SlotsShuffle(
    const std::set<std::string>& slots_to_replace) {
1744 1745
  PADDLE_ENFORCE_EQ(slots_shuffle_fea_eval_,
                    true,
1746 1747 1748 1749 1750 1751 1752
                    platform::errors::PreconditionNotMet(
                        "fea eval mode off, need to set on for slots shuffle"));
  platform::Timer timeline;
  timeline.Start();
  std::unordered_set<uint16_t> index_slots;
  PreprocessChannel(slots_to_replace, index_slots);

1753 1754 1755 1756 1757 1758 1759 1760 1761
  std::vector<Record> random_data;
  random_data.clear();
  // get slots shuffled random_data
  GetRandomData(index_slots, &random_data);
  input_channel_->Open();
  input_channel_->Write(std::move(random_data));
  random_data.clear();
  random_data.shrink_to_fit();
  input_channel_->Close();
Y
yaoxuefeng 已提交
1762
  cur_channel_ = 0;
1763 1764 1765 1766 1767 1768 1769

  timeline.Pause();
  VLOG(2) << "DatasetImpl<T>::SlotsShuffle() end"
          << ", memory data size for slots shuffle=" << input_channel_->Size()
          << ", cost time=" << timeline.ElapsedSec() << " seconds";
}

Y
yaoxuefeng 已提交
1770 1771 1772 1773 1774 1775 1776 1777 1778 1779 1780 1781 1782 1783 1784
template class DatasetImpl<SlotRecord>;
void SlotRecordDataset::CreateChannel() {
  if (input_channel_ == nullptr) {
    input_channel_ = paddle::framework::MakeChannel<SlotRecord>();
  }
}
void SlotRecordDataset::CreateReaders() {
  VLOG(3) << "Calling CreateReaders()";
  VLOG(3) << "thread num in Dataset: " << thread_num_;
  VLOG(3) << "Filelist size in Dataset: " << filelist_.size();
  VLOG(3) << "channel num in Dataset: " << channel_num_;
  CHECK(thread_num_ > 0) << "thread num should > 0";
  CHECK(channel_num_ > 0) << "channel num should > 0";
  CHECK(channel_num_ <= thread_num_) << "channel num should <= thread num";
  VLOG(3) << "readers size: " << readers_.size();
1785
  if (!readers_.empty()) {
Y
yaoxuefeng 已提交
1786 1787 1788 1789 1790 1791 1792 1793 1794 1795 1796 1797 1798 1799 1800 1801 1802 1803 1804 1805
    VLOG(3) << "readers_.size() = " << readers_.size()
            << ", will not create again";
    return;
  }
  VLOG(3) << "data feed class name: " << data_feed_desc_.name();
  for (int i = 0; i < thread_num_; ++i) {
    readers_.push_back(DataFeedFactory::CreateDataFeed(data_feed_desc_.name()));
    readers_[i]->Init(data_feed_desc_);
    readers_[i]->SetThreadId(i);
    readers_[i]->SetThreadNum(thread_num_);
    readers_[i]->SetFileListMutex(&mutex_for_pick_file_);
    readers_[i]->SetFileListIndex(&file_idx_);
    readers_[i]->SetFeaNumMutex(&mutex_for_fea_num_);
    readers_[i]->SetFeaNum(&total_fea_num_);
    readers_[i]->SetFileList(filelist_);
    readers_[i]->SetParseInsId(parse_ins_id_);
    readers_[i]->SetParseContent(parse_content_);
    readers_[i]->SetParseLogKey(parse_logkey_);
    readers_[i]->SetEnablePvMerge(enable_pv_merge_);
    readers_[i]->SetCurrentPhase(current_phase_);
L
lxsbupt 已提交
1806 1807 1808
#if defined(PADDLE_WITH_GPU_GRAPH) && defined(PADDLE_WITH_HETERPS)
    readers_[i]->InitGraphResource();
#endif
Y
yaoxuefeng 已提交
1809 1810 1811 1812 1813 1814 1815 1816 1817 1818 1819 1820 1821 1822 1823 1824 1825 1826 1827 1828 1829 1830 1831 1832 1833 1834 1835 1836 1837 1838 1839 1840 1841 1842 1843 1844 1845 1846 1847 1848 1849 1850 1851 1852 1853 1854 1855 1856 1857 1858 1859 1860 1861 1862 1863 1864 1865 1866 1867 1868 1869 1870 1871
    if (input_channel_ != nullptr) {
      readers_[i]->SetInputChannel(input_channel_.get());
    }
  }
  VLOG(3) << "readers size: " << readers_.size();
}

void SlotRecordDataset::ReleaseMemory() {
  VLOG(3) << "SlotRecordDataset::ReleaseMemory() begin";
  platform::Timer timeline;
  timeline.Start();

  if (input_channel_) {
    input_channel_->Clear();
    input_channel_ = nullptr;
  }
  if (enable_heterps_) {
    VLOG(3) << "put pool records size: " << input_records_.size();
    SlotRecordPool().put(&input_records_);
    input_records_.clear();
    input_records_.shrink_to_fit();
    VLOG(3) << "release heterps input records records size: "
            << input_records_.size();
  }

  readers_.clear();
  readers_.shrink_to_fit();

  std::vector<std::shared_ptr<paddle::framework::DataFeed>>().swap(readers_);

  VLOG(3) << "SlotRecordDataset::ReleaseMemory() end";
  VLOG(3) << "total_feasign_num_(" << STAT_GET(STAT_total_feasign_num_in_mem)
          << ") - current_fea_num_(" << total_fea_num_ << ") = ("
          << STAT_GET(STAT_total_feasign_num_in_mem) - total_fea_num_ << ")"
          << " object pool size=" << SlotRecordPool().capacity();  // For Debug
  STAT_SUB(STAT_total_feasign_num_in_mem, total_fea_num_);
}
void SlotRecordDataset::GlobalShuffle(int thread_num) {
  // TODO(yaoxuefeng)
  return;
}

void SlotRecordDataset::DynamicAdjustChannelNum(int channel_num,
                                                bool discard_remaining_ins) {
  if (channel_num_ == channel_num) {
    VLOG(3) << "DatasetImpl<T>::DynamicAdjustChannelNum channel_num_="
            << channel_num_ << ", channel_num_=channel_num, no need to adjust";
    return;
  }
  VLOG(3) << "adjust channel num from " << channel_num_ << " to "
          << channel_num;
  channel_num_ = channel_num;

  if (static_cast<int>(input_channel_->Size()) >= channel_num) {
    input_channel_->SetBlockSize(input_channel_->Size() / channel_num +
                                 (discard_remaining_ins ? 0 : 1));
  }

  VLOG(3) << "adjust channel num done";
}

void SlotRecordDataset::PrepareTrain() {
#ifdef PADDLE_WITH_GLOO
Y
yaoxuefeng 已提交
1872
  if (enable_heterps_) {
1873
    if (input_records_.empty() && input_channel_ != nullptr &&
Y
yaoxuefeng 已提交
1874 1875 1876 1877 1878 1879 1880 1881 1882 1883 1884 1885 1886 1887
        input_channel_->Size() != 0) {
      input_channel_->ReadAll(input_records_);
      VLOG(3) << "read from channel to records with records size: "
              << input_records_.size();
    }
    VLOG(3) << "input records size: " << input_records_.size();
    int64_t total_ins_num = input_records_.size();
    std::vector<std::pair<int, int>> offset;
    int default_batch_size =
        reinterpret_cast<SlotRecordInMemoryDataFeed*>(readers_[0].get())
            ->GetDefaultBatchSize();
    VLOG(3) << "thread_num: " << thread_num_
            << " memory size: " << total_ins_num
            << " default batch_size: " << default_batch_size;
1888 1889
    compute_thread_batch_nccl(
        thread_num_, total_ins_num, default_batch_size, &offset);
Y
yaoxuefeng 已提交
1890 1891 1892 1893 1894 1895 1896 1897 1898 1899 1900
    VLOG(3) << "offset size: " << offset.size();
    for (int i = 0; i < thread_num_; i++) {
      reinterpret_cast<SlotRecordInMemoryDataFeed*>(readers_[i].get())
          ->SetRecord(&input_records_[0]);
    }
    for (size_t i = 0; i < offset.size(); i++) {
      reinterpret_cast<SlotRecordInMemoryDataFeed*>(
          readers_[i % thread_num_].get())
          ->AddBatchOffset(offset[i]);
    }
  }
Y
yaoxuefeng 已提交
1901 1902 1903 1904 1905 1906 1907 1908 1909 1910 1911 1912 1913 1914 1915 1916 1917 1918 1919 1920 1921
#else
  PADDLE_THROW(platform::errors::Unavailable(
      "dataset set heterps need compile with GLOO"));
#endif
  return;
}

void SlotRecordDataset::DynamicAdjustReadersNum(int thread_num) {
  if (thread_num_ == thread_num) {
    VLOG(3) << "DatasetImpl<T>::DynamicAdjustReadersNum thread_num_="
            << thread_num_ << ", thread_num_=thread_num, no need to adjust";
    return;
  }
  VLOG(3) << "adjust readers num from " << thread_num_ << " to " << thread_num;
  thread_num_ = thread_num;
  std::vector<std::shared_ptr<paddle::framework::DataFeed>>().swap(readers_);
  CreateReaders();
  VLOG(3) << "adjust readers num done";
  PrepareTrain();
}

D
dongdaxiang 已提交
1922 1923
}  // end namespace framework
}  // end namespace paddle