| 45 | recent_request_ids_(100000) {} |
| 46 | |
| 47 | void GdrWorker::GrpcRecvTensorAsync(CallOptions* opts, |
| 48 | const RecvTensorRequest* request, |
| 49 | ::grpc::ByteBuffer* response, |
| 50 | StatusCallback done) { |
| 51 | Status s = recent_request_ids_.TrackUnique( |
| 52 | request->request_id(), "RecvTensor (GdrWorker)", *request); |
| 53 | if (!s.ok()) { |
| 54 | done(s); |
| 55 | return; |
| 56 | } |
| 57 | |
| 58 | const int64 step_id = request->step_id(); |
| 59 | const string& key = request->rendezvous_key(); |
| 60 | TRACEPRINTF("RecvTensor: %lld %s", step_id, key.c_str()); |
| 61 | Rendezvous::ParsedKey parsed; |
| 62 | s = Rendezvous::ParseKey(key, &parsed); |
| 63 | Device* src_dev = nullptr; |
| 64 | if (s.ok()) { |
| 65 | s = PrepareRecvTensor(parsed, &src_dev); |
| 66 | } |
| 67 | if (!s.ok()) { |
| 68 | done(s); |
| 69 | return; |
| 70 | } |
| 71 | |
| 72 | // Request the tensor associated with the rendezvous key. Any time |
| 73 | // while waiting for the tensor to be produced, up until the start |
| 74 | // of execution of the callback lambda body below, an RPC |
| 75 | // cancellation should abort the rendezvous. |
| 76 | opts->SetCancelCallback([this, step_id]() { AbortStep(step_id); }); |
| 77 | const bool dma_ok = request->dma_ok(); |
| 78 | env_->rendezvous_mgr->RecvLocalAsync( |
| 79 | step_id, parsed, |
| 80 | [this, opts, response, done, src_dev, request, dma_ok]( |
| 81 | const Status& status, const Rendezvous::Args& send_args, |
| 82 | const Rendezvous::Args&, const Tensor& val, const bool is_dead) { |
| 83 | opts->ClearCancelCallback(); |
| 84 | if (status.ok()) { |
| 85 | // DMA can only be used for Tensors that do not fall into |
| 86 | // the following three odd edge cases: 1) a zero-size |
| 87 | // buffer, 2) a dead tensor which has an uninit value, and |
| 88 | // 3) the tensor has the on_host allocation attribute, |
| 89 | // i.e. it's in CPU RAM *independent of its assigned |
| 90 | // device type*. |
| 91 | const bool on_host = send_args.alloc_attrs.on_host(); |
| 92 | if (val.TotalBytes() > 1024 && (!is_dead) && |
| 93 | DMAHelper::CanUseDMA(&val) && dma_ok) { |
| 94 | // DMA cases. |
| 95 | RecvTensorResponse* proto = new RecvTensorResponse; |
| 96 | proto->set_is_dead(is_dead); |
| 97 | proto->set_send_start_micros(Env::Default()->NowMicros()); |
| 98 | TensorProto* tensor_proto = proto->mutable_tensor(); |
| 99 | tensor_proto->set_dtype(val.dtype()); |
| 100 | val.shape().AsProto(tensor_proto->mutable_tensor_shape()); |
| 101 | auto transport_options = proto->mutable_transport_options(); |
| 102 | remote_memory_manager_->TransportOptionsFromTensor( |
| 103 | transport_options, val, src_dev, send_args.device_context, |
| 104 | on_host, [proto, done, response](const Status& s) { |
nothing calls this directly
no test coverage detected