| 342 | } |
| 343 | |
| 344 | Status SliceRecvOp::RecvTotalBytes(OpKernelContext* ctx, |
| 345 | const FrameAndIter& frame_iter, |
| 346 | bool& is_dead, uint64& total_bytes) { |
| 347 | Rendezvous::Args args; |
| 348 | args.device_context = ctx->op_device_context(); |
| 349 | args.alloc_attrs = AllocatorAttributes(); |
| 350 | if (ctx->is_eager()) { |
| 351 | // NOTE(fishx): Only set cancellation_manager in eager mode. Because in |
| 352 | // Tensorflow 1.x, session (or graph_mgr) will abort the underlying |
| 353 | // rendezvous if it encounters any error. |
| 354 | args.cancellation_manager = ctx->cancellation_manager(); |
| 355 | } |
| 356 | |
| 357 | Rendezvous::ParsedKey parsed_key; |
| 358 | Tensor total_bytes_t; |
| 359 | slice_sendrecv::GetSliceRendezvousKey(key_prefix_, |
| 360 | "_slice_transfer_totalbytes", frame_iter, &parsed_key.buf_); |
| 361 | VLOG(2) << "SliceRecv " << parsed_key.buf_; |
| 362 | TF_RETURN_IF_ERROR(Rendezvous::ParseKey(parsed_key.buf_, &parsed_key)); |
| 363 | TF_RETURN_IF_ERROR(ctx->rendezvous()->Recv(parsed_key, args, &total_bytes_t, |
| 364 | &is_dead, timeout_ms_)); |
| 365 | if (!is_dead) { |
| 366 | total_bytes = total_bytes_t.scalar<uint64>()(); |
| 367 | } |
| 368 | |
| 369 | return Status::OK(); |
| 370 | } |
| 371 | |
| 372 | Status SliceRecvOp::RecvShape(OpKernelContext* ctx, |
| 373 | const FrameAndIter& frame_iter, |
nothing calls this directly
no test coverage detected