| 208 | |
| 209 | #if !defined(IS_MOBILE_PLATFORM) |
| 210 | void EagerContext::CloseRemoteContexts() { |
| 211 | // Close all remote contexts. |
| 212 | eager::CloseContextRequest request; |
| 213 | uint64 context_id; |
| 214 | { |
| 215 | mutex_lock l(remote_state_mu_); |
| 216 | if (!is_master_) return; |
| 217 | context_id = context_id_; |
| 218 | context_id_ = kInvalidContextId; |
| 219 | } |
| 220 | request.set_context_id(context_id); |
| 221 | // Setting context_id to a new value can avoid us issuing DestroyTensorHandle |
| 222 | // request to closed remote workers. |
| 223 | std::vector<eager::CloseContextResponse> responses(remote_contexts_.size()); |
| 224 | BlockingCounter counter(static_cast<int>(remote_contexts_.size())); |
| 225 | |
| 226 | int i = 0; |
| 227 | for (const auto& worker : remote_contexts_) { |
| 228 | eager::EagerClient* client; |
| 229 | Status s = remote_eager_workers_->GetClient(worker, &client); |
| 230 | |
| 231 | client->CloseContextAsync( |
| 232 | &request, &responses[i], |
| 233 | [&worker, &counter, context_id](const Status& s) { |
| 234 | if (!s.ok()) { |
| 235 | LOG(ERROR) << "Unable to close remote context with ID " |
| 236 | << context_id << " for worker: " << worker << " due to " |
| 237 | << s.error_message(); |
| 238 | } |
| 239 | counter.DecrementCount(); |
| 240 | }); |
| 241 | i++; |
| 242 | } |
| 243 | |
| 244 | counter.Wait(); |
| 245 | |
| 246 | remote_contexts_.clear(); |
| 247 | } |
| 248 | |
| 249 | #endif // !IS_MOBILE_PLATFORM |
| 250 | |