| 419 | } |
| 420 | |
| 421 | void StarRemoteRendezvous::FuseRecvFromRemoteAsync( |
| 422 | const std::vector<Rendezvous::ParsedKey>& parsed_keys, |
| 423 | const Rendezvous::Args& recv_args, FuseDoneCallback done) { |
| 424 | CHECK(is_initialized()); |
| 425 | int fuse_count = parsed_keys.size(); |
| 426 | Status s; |
| 427 | |
| 428 | // Prepare a FuseRecvTensor call that can handle being aborted. |
| 429 | StarFuseRecvTensorCall* call = get_fuse_call_freelist()->New(); |
| 430 | |
| 431 | // key.src_device identifies a remote device. |
| 432 | if (!DeviceNameUtils::SplitDeviceName(parsed_keys[0].src_device, |
| 433 | &call->src_worker_, |
| 434 | &call->src_rel_device_)) { |
| 435 | s = errors::Internal(parsed_keys[0].src_device, |
| 436 | " is invalid remote source device."); |
| 437 | } |
| 438 | WorkerSession* sess = session(); |
| 439 | WorkerInterface* rwi = sess->worker_cache->GetOrCreateWorker(call->src_worker_); |
| 440 | if (s.ok() && rwi == nullptr) { |
| 441 | s = errors::Internal("No worker known as ", call->src_worker_); |
| 442 | } |
| 443 | |
| 444 | Device* dst_device; |
| 445 | if (s.ok()) { |
| 446 | s = sess->device_mgr()->LookupDevice(parsed_keys[0].dst_device, &dst_device); |
| 447 | } |
| 448 | if (!s.ok()) { |
| 449 | if (rwi != nullptr) { |
| 450 | sess->worker_cache->ReleaseWorker(call->src_worker_, rwi); |
| 451 | } |
| 452 | get_fuse_call_freelist()->Release(call, sess->worker_cache.get()); |
| 453 | done(s, std::vector<Args>(fuse_count), recv_args, |
| 454 | std::vector<Tensor>(fuse_count), |
| 455 | std::vector<bool>(fuse_count, false)); |
| 456 | return; |
| 457 | } |
| 458 | |
| 459 | call->Init(rwi, step_id_, |
| 460 | parsed_keys, recv_args.alloc_attrs, dst_device, |
| 461 | recv_args, std::move(done)); |
| 462 | |
| 463 | // Record "call" in active_ so that it can be aborted cleanly. |
| 464 | RegisterCall(call, recv_args); |
| 465 | Ref(); |
| 466 | if (!s.ok()) { |
| 467 | LOG(WARNING) << "Rendezvous has been aborted, ignore the rpc call." |
| 468 | << ", rendezvous key: " << parsed_keys[0].FullKey() |
| 469 | << ", step id: " << step_id_; |
| 470 | call->fuse_done()(s, std::vector<Args>(fuse_count), recv_args, |
| 471 | std::vector<Tensor>(fuse_count), |
| 472 | std::vector<bool>(fuse_count, false)); |
| 473 | session()->worker_cache->ReleaseWorker(call->src_worker_, call->wi_); |
| 474 | call->wi_ = nullptr; |
| 475 | get_fuse_call_freelist()->Release(call, session()->worker_cache.get()); |
| 476 | Unref(); |
| 477 | return; |
| 478 | } |
nothing calls this directly
no test coverage detected