| 93 | XrtGrpcEagerClientCache::~XrtGrpcEagerClientCache() { threads_.clear(); } |
| 94 | |
| 95 | xla::StatusOr<XrtGrpcEagerClient*> XrtGrpcEagerClientCache::GetClient( |
| 96 | const string& target) { |
| 97 | auto it = clients_.find(target); |
| 98 | if (it == clients_.end()) { |
| 99 | tensorflow::SharedGrpcChannelPtr shared = cache_->FindWorkerChannel(target); |
| 100 | if (!shared) { |
| 101 | return errors::NotFound("Unknown target ", target); |
| 102 | } |
| 103 | auto worker = absl::make_unique<XrtGrpcEagerClient>( |
| 104 | shared, threads_[AssignClientToThread(target)].completion_queue()); |
| 105 | |
| 106 | it = clients_.emplace(target, std::move(worker)).first; |
| 107 | } |
| 108 | |
| 109 | return it->second.get(); |
| 110 | } |
| 111 | |
| 112 | size_t XrtGrpcEagerClientCache::AssignClientToThread(const string& target) { |
| 113 | // Round-robin target assignment, but keeps the same target on the same |