| 410 | } |
| 411 | |
| 412 | void StarWorker::RecvTensorAsync(CallOptions* opts, |
| 413 | const RecvTensorRequest* request, |
| 414 | StarTensorResponse* response, |
| 415 | StatusCallback done) { |
| 416 | const int64 step_id = request->step_id(); |
| 417 | const string& key = request->rendezvous_key(); |
| 418 | Rendezvous::ParsedKey parsed; |
| 419 | |
| 420 | Status s = Rendezvous::ParseKey(key, &parsed); |
| 421 | Device* src_dev = nullptr; |
| 422 | if (s.ok()) { |
| 423 | s = PrepareRecvTensor(parsed, &src_dev); |
| 424 | } |
| 425 | if (!s.ok()) { |
| 426 | LOG(ERROR) << "PrepareRecvTensor failed, tensor:" << key; |
| 427 | done(s); |
| 428 | return; |
| 429 | } |
| 430 | |
| 431 | // TODO(rangeng.llb): make call opts useful. |
| 432 | // opts->SetCancelCallback([this, step_id]() { AbortStep(step_id); }); |
| 433 | env_->rendezvous_mgr->RecvLocalAsync( |
| 434 | step_id, parsed, |
| 435 | [opts, request, response, done, src_dev, key](const Status& status, |
| 436 | const Rendezvous::Args& send_args, |
| 437 | const Rendezvous::Args& recv_args, |
| 438 | const Tensor& val, const bool is_dead) { |
| 439 | //opts->ClearCancelCallback(); |
| 440 | |
| 441 | if (!status.ok()) { |
| 442 | LOG(ERROR) << "env_->rendezvous_mgr->RecvLocalAsync failed, error msg is: " |
| 443 | << status.error_message(); |
| 444 | } |
| 445 | |
| 446 | if (status.ok()) { |
| 447 | response->SetIsDead(is_dead); |
| 448 | bool can_memcpy = DataTypeCanUseMemcpy(val.dtype()); |
| 449 | |
| 450 | if (src_dev->tensorflow_gpu_device_info() && |
| 451 | (!send_args.alloc_attrs.on_host())) { |
| 452 | #if GOOGLE_CUDA |
| 453 | CHECK(send_args.device_context) |
| 454 | << "send dev name: " << src_dev->name() |
| 455 | << " gpu_info: " << src_dev->tensorflow_gpu_device_info(); |
| 456 | |
| 457 | if (can_memcpy) { |
| 458 | Allocator* alloc = GPUProcessState::singleton()->GetGpuHostAllocator(0); |
| 459 | Tensor* cpu_copy = new Tensor(alloc, val.dtype(), val.shape()); |
| 460 | |
| 461 | GPUUtil::CopyGPUTensorToCPU(src_dev, send_args.device_context, &val, cpu_copy, |
| 462 | [response, cpu_copy, done](const Status& s) { |
| 463 | CHECK(s.ok()) << "copy tensor from gpu sync"; |
| 464 | response->SetTensor(*cpu_copy); |
| 465 | delete cpu_copy; |
| 466 | done(s); |
| 467 | }); |
| 468 | } else { |
| 469 | // TODO(rangeng.llb): Should not be executed currrently. |
no test coverage detected