| 165 | } |
| 166 | |
| 167 | void EagerExecutor::Run() { |
| 168 | auto thread_exited_notifier = |
| 169 | gtl::MakeCleanup([this] { thread_exited_notification_.Notify(); }); |
| 170 | while (true) { |
| 171 | EagerNode* curr_node_raw; |
| 172 | { |
| 173 | tensorflow::mutex_lock l(node_queue_mutex_); |
| 174 | while (node_queue_.empty() || !status_.ok()) { |
| 175 | if (state_ == ExecutorState::kShutDown) return; |
| 176 | nodes_pending_.wait(l); |
| 177 | } |
| 178 | // Obtain raw pointer since we don't want to remove from the queue until |
| 179 | // the node has been run. Otherwise, WaitForAllPendingNodes can return |
| 180 | // too early. |
| 181 | // Note, we don't std::move from the here because the front of the queue |
| 182 | // will then contain a nullptr. This can be a problem in |
| 183 | // WaitForAllPendingNodes where we get the top EagerNode pointer |
| 184 | // and register a notification for its completion. |
| 185 | curr_node_raw = node_queue_.front().get(); |
| 186 | } |
| 187 | tensorflow::Status status = curr_node_raw->Run(); |
| 188 | const bool ok = status.ok(); |
| 189 | |
| 190 | std::unique_ptr<EagerNode> curr_node; |
| 191 | std::vector<std::unique_ptr<EagerNode>> nodes_to_destroy; |
| 192 | { |
| 193 | tensorflow::mutex_lock l(node_queue_mutex_); |
| 194 | curr_node = std::move(node_queue_.front()); |
| 195 | node_queue_.pop(); |
| 196 | if (!ok) { |
| 197 | status_ = status; |
| 198 | // We remove any pending ops so that we don't try to execute them if |
| 199 | // ClearError is called. |
| 200 | errors::AppendToMessage( |
| 201 | &status, |
| 202 | ". Encountered when executing an operation using " |
| 203 | "EagerExecutor. This error cancels all future " |
| 204 | "operations and poisons their output tensors."); |
| 205 | while (!node_queue_.empty()) { |
| 206 | node_queue_.front()->Abort(status); |
| 207 | nodes_to_destroy.push_back(std::move(node_queue_.front())); |
| 208 | node_queue_.pop(); |
| 209 | } |
| 210 | } |
| 211 | if (!node_done_notifications_.empty()) { |
| 212 | // Note that we notify all waiting threads in case an error has |
| 213 | // occurred. These calling threads are responsible for checking status_ |
| 214 | // before proceeding. |
| 215 | const auto range = |
| 216 | ok ? node_done_notifications_.equal_range(curr_node_raw) |
| 217 | : make_pair(node_done_notifications_.begin(), |
| 218 | node_done_notifications_.end()); |
| 219 | for (auto it = range.first; it != range.second; ++it) { |
| 220 | it->second->notify_all(); |
| 221 | } |
| 222 | node_done_notifications_.erase(range.first, range.second); |
| 223 | } |
| 224 | } |