| 45 | } |
| 46 | |
| 47 | void FileSliceSendOp::Compute(OpKernelContext* ctx) { |
| 48 | OP_REQUIRES(ctx, ctx->rendezvous() != nullptr, |
| 49 | errors::Internal("Op kernel context needs to provide a rendezvous.")); |
| 50 | |
| 51 | const Tensor& file_path_t = ctx->input(0); |
| 52 | if (!ctx->is_input_dead()) { |
| 53 | OP_REQUIRES(ctx, TensorShapeUtils::IsScalar(file_path_t.shape()), |
| 54 | errors::InvalidArgument("file_path is not a scalar: ", |
| 55 | file_path_t.shape().DebugString())); |
| 56 | } |
| 57 | |
| 58 | FrameAndIter frame_iter = \ |
| 59 | slice_sendrecv::GetFrameAndIter(ctx, hostmem_sendrecv_); |
| 60 | |
| 61 | // get element_bytes. |
| 62 | uint64 element_bytes = 0; |
| 63 | OP_REQUIRES_OK(ctx, GetElementBytes(ctx, file_path_t, element_bytes)); |
| 64 | |
| 65 | // send total_bytes. |
| 66 | // total_bytes is the TotalBytes of the Tensor that contains the contents of |
| 67 | // the file. please refer Tensor::TotalBytes() |
| 68 | uint64 total_bytes = element_bytes + sizeof(tstring); |
| 69 | OP_REQUIRES_OK(ctx, SendTotalBytes(ctx, frame_iter, total_bytes)); |
| 70 | // if input is dead, only send total_bytes dead tensor. |
| 71 | if (ctx->is_input_dead()) { |
| 72 | return; |
| 73 | } |
| 74 | |
| 75 | // if total bytes is smaller than slice size, send directly. |
| 76 | if (total_bytes <= slice_size_) { |
| 77 | Rendezvous::Args args; |
| 78 | args.device_context = ctx->op_device_context(); |
| 79 | args.alloc_attrs = ctx->input_alloc_attr(0); |
| 80 | |
| 81 | Rendezvous::ParsedKey parsed_key; |
| 82 | slice_sendrecv::GetSliceRendezvousKey(key_prefix_, "_transfer_data", |
| 83 | frame_iter, &parsed_key.buf_); |
| 84 | VLOG(2) << "FileSliceSend " << parsed_key.buf_; |
| 85 | OP_REQUIRES_OK(ctx, Rendezvous::ParseKey(parsed_key.buf_, &parsed_key)); |
| 86 | Tensor data_t; |
| 87 | OP_REQUIRES_OK(ctx, |
| 88 | ctx->allocate_temp(DT_STRING, TensorShape({}), &data_t)); |
| 89 | if (element_bytes > 0) { |
| 90 | OP_REQUIRES_OK(ctx, ReadFileToString(Env::Default(), |
| 91 | file_path_t.scalar<tstring>()(), data_t.scalar<tstring>().data())); |
| 92 | } |
| 93 | OP_REQUIRES_OK(ctx, ctx->rendezvous()->Send(parsed_key,args, data_t, |
| 94 | ctx->is_input_dead())); |
| 95 | return; |
| 96 | } |
| 97 | |
| 98 | // send shape, in order to match the behavior of 'SliceSend'. |
| 99 | OP_REQUIRES_OK(ctx, SendScalarShape(ctx, frame_iter)); |
| 100 | |
| 101 | // send element bytes, in order to match the behavior of 'SliceSend'. |
| 102 | OP_REQUIRES_OK(ctx, SendElementBytes(ctx, frame_iter, element_bytes)); |
| 103 | |
| 104 | // send data. |
nothing calls this directly
no test coverage detected