| 399 | } |
| 400 | |
| 401 | Status FileSliceRecvOp::RecvFile(OpKernelContext* ctx, |
| 402 | const FrameAndIter& frame_iter, |
| 403 | const string& file_path) { |
| 404 | Rendezvous::Args args; |
| 405 | args.device_context = ctx->op_device_context(); |
| 406 | args.alloc_attrs = ctx->output_alloc_attr(0); |
| 407 | if (ctx->is_eager()) { |
| 408 | // NOTE(fishx): Only set cancellation_manager in eager mode. Because in |
| 409 | // Tensorflow 1.x, session (or graph_mgr) will abort the underlying |
| 410 | // rendezvous if it encounters any error. |
| 411 | args.cancellation_manager = ctx->cancellation_manager(); |
| 412 | } |
| 413 | |
| 414 | Rendezvous::ParsedKey parsed_key; |
| 415 | slice_sendrecv::GetSliceRendezvousKey(key_prefix_, "_transfer_data", |
| 416 | frame_iter, &parsed_key.buf_); |
| 417 | VLOG(2) << "FileSliceRecv " << parsed_key.buf_; |
| 418 | TF_RETURN_IF_ERROR(Rendezvous::ParseKey(parsed_key.buf_, &parsed_key)); |
| 419 | Tensor data_t; |
| 420 | bool is_dead = false; |
| 421 | TF_RETURN_IF_ERROR(ctx->rendezvous()->Recv(parsed_key, args, &data_t, |
| 422 | &is_dead, timeout_ms_)); |
| 423 | |
| 424 | // This shouldn't be a dead tensor. |
| 425 | CHECK_EQ(is_dead, false); |
| 426 | |
| 427 | // Write data_t to file. |
| 428 | Env* env = Env::Default(); |
| 429 | return WriteStringToFile(env, file_path, data_t.scalar<tstring>()()); |
| 430 | } |
| 431 | |
| 432 | Status FileSliceRecvOp::RecvFileSlice(OpKernelContext* ctx, |
| 433 | const FrameAndIter& frame_iter, |
nothing calls this directly
no test coverage detected