| 65 | } |
| 66 | |
| 67 | void PartialRunMgr::PartialRunDone(int step_id, StatusCallback done, |
| 68 | const Status& status) { |
| 69 | Status callback_status; |
| 70 | { |
| 71 | mutex_lock l(mu_); |
| 72 | auto run_it = step_id_to_partial_run_.find(step_id); |
| 73 | if (run_it == step_id_to_partial_run_.end()) { |
| 74 | return; |
| 75 | } |
| 76 | run_it->second->final_status.Update(status); |
| 77 | if (!run_it->second->executor_done) { |
| 78 | // If we found the partial_run, we set the final callback to call only |
| 79 | // when the executor is completely done. |
| 80 | run_it->second->final_callback = std::move(done); |
| 81 | return; |
| 82 | } |
| 83 | callback_status = run_it->second->final_status; |
| 84 | } |
| 85 | // Otherwise we call the callback immediately. |
| 86 | done(callback_status); |
| 87 | mutex_lock l(mu_); |
| 88 | step_id_to_partial_run_.erase(step_id); |
| 89 | } |
| 90 | |
| 91 | } // namespace tensorflow |