| 457 | } |
| 458 | |
| 459 | Status SliceRecvOp::RecvStringSlice(OpKernelContext* ctx, |
| 460 | const FrameAndIter& frame_iter, |
| 461 | const int64 index, |
| 462 | const uint64 element_bytes, |
| 463 | TTypes<tstring>::Flat& output_flat) { |
| 464 | Rendezvous::Args args; |
| 465 | args.device_context = ctx->op_device_context(); |
| 466 | args.alloc_attrs = ctx->output_alloc_attr(0); |
| 467 | if (ctx->is_eager()) { |
| 468 | // NOTE(fishx): Only set cancellation_manager in eager mode. Because in |
| 469 | // Tensorflow 1.x, session (or graph_mgr) will abort the underlying |
| 470 | // rendezvous if it encounters any error. |
| 471 | args.cancellation_manager = ctx->cancellation_manager(); |
| 472 | } |
| 473 | Rendezvous::ParsedKey parsed_key; |
| 474 | |
| 475 | int64 slice_num = element_bytes / slice_size_; |
| 476 | if (element_bytes % slice_size_ != 0) { |
| 477 | slice_num += 1; |
| 478 | } |
| 479 | Tensor data_t; |
| 480 | bool is_dead = false; |
| 481 | for (int64 i = 0; i < slice_num; i++) { |
| 482 | std::string tensor_name_suffix = \ |
| 483 | strings::StrCat("_slice_transfer_data_", std::to_string(index), "_", |
| 484 | std::to_string(i)); |
| 485 | slice_sendrecv::GetSliceRendezvousKey(key_prefix_, tensor_name_suffix, |
| 486 | frame_iter, &parsed_key.buf_); |
| 487 | VLOG(2) << "SliceRecv " << parsed_key.buf_; |
| 488 | TF_RETURN_IF_ERROR(Rendezvous::ParseKey(parsed_key.buf_, &parsed_key)); |
| 489 | TF_RETURN_IF_ERROR( |
| 490 | ctx->rendezvous()->FlowControlRecv(tensor_name_, parsed_key, args, |
| 491 | &data_t, &is_dead, timeout_ms_)); |
| 492 | // This shouldn't be a dead tensor. |
| 493 | CHECK_EQ(is_dead, false); |
| 494 | output_flat(index) += data_t.scalar<tstring>()(); |
| 495 | } |
| 496 | |
| 497 | return Status::OK(); |
| 498 | } |
| 499 | |
| 500 | Status SliceRecvOp::RecvBasicType(OpKernelContext* ctx, |
| 501 | const FrameAndIter& frame_iter, |
nothing calls this directly
no test coverage detected