| 404 | } |
| 405 | |
| 406 | Status SliceRecvOp::RecvString(OpKernelContext* ctx, |
| 407 | const FrameAndIter& frame_iter, |
| 408 | const TensorShape& shape, Tensor*& output_t) { |
| 409 | Rendezvous::Args args; |
| 410 | args.device_context = ctx->op_device_context(); |
| 411 | args.alloc_attrs = AllocatorAttributes(); |
| 412 | if (ctx->is_eager()) { |
| 413 | // NOTE(fishx): Only set cancellation_manager in eager mode. Because in |
| 414 | // Tensorflow 1.x, session (or graph_mgr) will abort the underlying |
| 415 | // rendezvous if it encounters any error. |
| 416 | args.cancellation_manager = ctx->cancellation_manager(); |
| 417 | } |
| 418 | Rendezvous::ParsedKey parsed_key; |
| 419 | bool is_dead; |
| 420 | |
| 421 | // recv elements bytes. |
| 422 | slice_sendrecv::GetSliceRendezvousKey(key_prefix_, |
| 423 | "_slice_transfer_elements_bytes", frame_iter, &parsed_key.buf_); |
| 424 | VLOG(2) << "SliceRecv " << parsed_key.buf_; |
| 425 | TF_RETURN_IF_ERROR(Rendezvous::ParseKey(parsed_key.buf_, &parsed_key)); |
| 426 | Tensor elements_bytes_t; |
| 427 | TF_RETURN_IF_ERROR(ctx->rendezvous()->Recv(parsed_key, args, &elements_bytes_t, |
| 428 | &is_dead, timeout_ms_)); |
| 429 | // This shouldn't be a dead tensor. |
| 430 | CHECK_EQ(is_dead, false); |
| 431 | auto elements_bytes_flat = elements_bytes_t.flat<uint64>(); |
| 432 | int64 num_elements = shape.num_elements(); |
| 433 | args.alloc_attrs = ctx->output_alloc_attr(0); |
| 434 | Tensor data_t; |
| 435 | auto output_flat = output_t->flat<tstring>(); |
| 436 | for (int64 i = 0; i < num_elements; i++) { |
| 437 | if (elements_bytes_flat(i) <= slice_size_) { |
| 438 | std::string tensor_name_suffix = \ |
| 439 | strings::StrCat("_slice_transfer_data_", std::to_string(i)); |
| 440 | slice_sendrecv::GetSliceRendezvousKey(key_prefix_, tensor_name_suffix, |
| 441 | frame_iter, &parsed_key.buf_); |
| 442 | VLOG(2) << "SliceRecv " << parsed_key.buf_; |
| 443 | TF_RETURN_IF_ERROR(Rendezvous::ParseKey(parsed_key.buf_, &parsed_key)); |
| 444 | TF_RETURN_IF_ERROR( |
| 445 | ctx->rendezvous()->FlowControlRecv(tensor_name_, parsed_key, args, |
| 446 | &data_t, &is_dead, timeout_ms_)); |
| 447 | // This shouldn't be a dead tensor. |
| 448 | CHECK_EQ(is_dead, false); |
| 449 | output_flat(i) = data_t.scalar<tstring>()(); |
| 450 | } else { |
| 451 | TF_RETURN_IF_ERROR(RecvStringSlice(ctx, frame_iter, i, |
| 452 | elements_bytes_flat(i), output_flat)); |
| 453 | } |
| 454 | } |
| 455 | |
| 456 | return Status::OK(); |
| 457 | } |
| 458 | |
| 459 | Status SliceRecvOp::RecvStringSlice(OpKernelContext* ctx, |
| 460 | const FrameAndIter& frame_iter, |
nothing calls this directly
no test coverage detected