Method RecvFileSlice

tensorflow/core/kernels/file_slice_sendrecv_ops.cc:432–475

}

Status FileSliceRecvOp::RecvFileSlice(OpKernelContext* ctx,
                                      const FrameAndIter& frame_iter,
                                      const uint64 element_bytes,
                                      const string& file_path) {
  // create file
  Env* env = Env::Default();
  std::unique_ptr<WritableFile> file_ptr;
  TF_RETURN_IF_ERROR(env->NewWritableFile(file_path, &file_ptr));

  Rendezvous::Args args;
  args.device_context = ctx->op_device_context();
  args.alloc_attrs = ctx->output_alloc_attr(0);
  if (ctx->is_eager()) {
    // NOTE(fishx): Only set cancellation_manager in eager mode. Because in
    // Tensorflow 1.x, session (or graph_mgr) will abort the underlying
    // rendezvous if it encounters any error.
    args.cancellation_manager = ctx->cancellation_manager();
  }
  Rendezvous::ParsedKey parsed_key;

  int64 slice_num = element_bytes / slice_size_;
  if (element_bytes % slice_size_ != 0) {
    slice_num += 1;
  }
  Tensor data_t;
  bool is_dead = false;
  for (int64 i = 0; i < slice_num; i++) {
    std::string tensor_name_suffix = \
      strings::StrCat("_slice_transfer_data_", std::to_string(0), "_",
                      std::to_string(i));
    slice_sendrecv::GetSliceRendezvousKey(key_prefix_, tensor_name_suffix,
                                          frame_iter, &parsed_key.buf_);
    VLOG(2) << "FileSliceRecv " << parsed_key.buf_;
    TF_RETURN_IF_ERROR(Rendezvous::ParseKey(parsed_key.buf_, &parsed_key));
    TF_RETURN_IF_ERROR(
      ctx->rendezvous()->FlowControlRecv(tensor_name_, parsed_key, args,
                                         &data_t, &is_dead, timeout_ms_));
    // This shouldn't be a dead tensor.
    CHECK_EQ(is_dead, false);
    file_ptr->Append(data_t.scalar<tstring>()());
  }

  return Status::OK();
}

REGISTER_KERNEL_BUILDER(Name("_FileSliceRecv").Device(DEVICE_CPU),
                        FileSliceRecvOp);