Method MainLoop

tensorflow/core/kernels/record_yielder.cc:97–172

}

void RecordYielder::MainLoop() {
  while (true) {
    ++epoch_;
    num_records_yielded_in_epoch_ = 0;
    num_records_added_in_epoch_ = 0;

    // Finds all files.
    std::vector<string> filenames;
    Status s = MatchFiles(opts_.file_pattern, &filenames);

    if (filenames.empty()) {
      s = errors::NotFound("Found no files at ", opts_.file_pattern);
      if (ShouldFinish(s)) {
        buf_enough_.notify_all();
        break;
      }
    }

    if (ShouldFinish(s)) break;

    // Shuffles these files according to the epoch # and random seed.
    std::mt19937_64 shuffle_rnd(
        Hash64(reinterpret_cast<char*>(&epoch_), sizeof(epoch_), opts_.seed));
    std::shuffle(filenames.begin(), filenames.end(), shuffle_rnd);

    // Left-shift the filename list.
    const std::vector<string>::size_type num = filenames.size();
    int64 shift;
    if (0 <= opts_.file_shuffle_shift_ratio &&
        opts_.file_shuffle_shift_ratio < 1) {
      shift = opts_.file_shuffle_shift_ratio * num;
      std::rotate(filenames.begin(), filenames.begin() + shift,
                  filenames.end());
    }

    // Shards files and use one thread to go through each shard.
    const int N = opts_.parallelism;
    std::vector<Shard> shards(N);
    for (int i = 0; i < N; ++i) {
      Shard* shard = &shards[i];
      shard->index = i;
      for (std::vector<string>::size_type j = i; j < filenames.size(); j += N) {
        shard->filenames.push_back(filenames[j]);
      }
      thread_->Schedule([this, shard]() { ShardLoop(shard); });
    }
    for (int i = 0; i < N; ++i) {
      shards[i].done.WaitForNotification();
      s.Update(shards[i].status);
    }

    if (num_records_added_in_epoch_ < opts_.bufsize) {
      mutex_lock l(mu_);
      opts_.bufsize = num_records_added_in_epoch_;
    }

    if (ShouldFinish(s)) {
      buf_enough_.notify_all();