| 2429 | } |
| 2430 | |
| 2431 | void DirectSession::WaitForNotification(RunState* run_state, |
| 2432 | CancellationManager* cm, |
| 2433 | int64 timeout_in_ms) { |
| 2434 | const Status status = |
| 2435 | WaitForNotification(&run_state->executors_done, timeout_in_ms); |
| 2436 | if (!status.ok()) { |
| 2437 | { |
| 2438 | mutex_lock l(run_state->mu_); |
| 2439 | run_state->status.Update(status); |
| 2440 | } |
| 2441 | cm->StartCancel(); |
| 2442 | // We must wait for the executors to complete, because they have borrowed |
| 2443 | // references to `cm` and other per-step state. After this notification, it |
| 2444 | // is safe to clean up the step. |
| 2445 | run_state->executors_done.WaitForNotification(); |
| 2446 | } |
| 2447 | } |
| 2448 | |
| 2449 | ::tensorflow::Status DirectSession::WaitForNotification( |
| 2450 | Notification* notification, int64 timeout_in_ms) { |