| 98 | } |
| 99 | |
| 100 | Status RemoteCopyNode::StartSend() { |
| 101 | // TODO(gjn): We should consider just using the low-level SendOp::Compute() |
| 102 | // functionality here instead of constructing an Op. |
| 103 | const AttrTypeMap* types; |
| 104 | bool is_function = false; |
| 105 | Status status = AttrTypeMapForOp("_Send", &types, &is_function); |
| 106 | if (!status.ok()) { |
| 107 | captured_state_->SetSendStatus(status); |
| 108 | return status; |
| 109 | } |
| 110 | DCHECK(!is_function); |
| 111 | EagerOperation op(ctx_, "_Send", /*is_function=*/false, types); |
| 112 | |
| 113 | op.SetDevice(send_device_); |
| 114 | |
| 115 | op.MutableAttrs()->Set("tensor_name", wire_id_); |
| 116 | op.MutableAttrs()->Set("send_device", send_device_->name()); |
| 117 | op.MutableAttrs()->Set( |
| 118 | "send_device_incarnation", |
| 119 | static_cast<int64>(send_device_->attributes().incarnation())); |
| 120 | op.MutableAttrs()->Set("recv_device", recv_device_->name()); |
| 121 | op.MutableAttrs()->Set("client_terminated", false); |
| 122 | |
| 123 | op.MutableAttrs()->Set("T", src_->dtype); |
| 124 | |
| 125 | DCHECK(send_device_ != nullptr); |
| 126 | |
| 127 | if (send_device_->IsLocal()) { |
| 128 | status = RunLocalSend(&op); |
| 129 | captured_state_->SetSendStatus(status); |
| 130 | return status; |
| 131 | } else { |
| 132 | // Prepare the request |
| 133 | EnqueueRequest request; |
| 134 | request.set_context_id(ctx_->GetContextId()); |
| 135 | auto* remote_op = request.add_queue()->mutable_operation(); |
| 136 | status = ctx_->RemoteMgr()->SerializeRemoteTensorHandle( |
| 137 | src_, remote_op->add_inputs(), src_->device(), |
| 138 | src_->DeviceOrHostCPU(ctx_)->name()); |
| 139 | if (!status.ok()) { |
| 140 | captured_state_->SetSendStatus(status); |
| 141 | return status; |
| 142 | } |
| 143 | |
| 144 | PrepareRemoteOp(remote_op, &op); |
| 145 | remote_op->set_id(ctx_->RemoteMgr()->NextOpId()); |
| 146 | |
| 147 | // Issue the RPC |
| 148 | eager::EagerClient* eager_client; |
| 149 | status = ctx_->GetClient(send_device_, &eager_client); |
| 150 | if (!status.ok()) { |
| 151 | captured_state_->SetSendStatus(status); |
| 152 | return status; |
| 153 | } |
| 154 | |
| 155 | const std::shared_ptr<CapturedSharedState>& captured_state = |
| 156 | captured_state_; |
| 157 | EnqueueResponse* response = new EnqueueResponse; |
nothing calls this directly
no test coverage detected