| 67 | step_stats_collector_(step_stats_collector) {} |
| 68 | |
| 69 | void NodeExecStatsWrapper::Done(const string& device) { |
| 70 | // TODO(tucker): merge with the DetailText function in session.cc in a common |
| 71 | // location. |
| 72 | DCHECK(node_); |
| 73 | string memory; |
| 74 | for (auto& all : stats_->memory()) { |
| 75 | int64 tot = all.total_bytes(); |
| 76 | if (tot >= 0.1 * 1048576.0) { |
| 77 | int64 peak = all.peak_bytes(); |
| 78 | if (peak > 0) { |
| 79 | memory = |
| 80 | strings::StrCat(memory, "[", all.allocator_name(), |
| 81 | strings::Printf(" %.1fMB %.1fMB] ", tot / 1048576.0, |
| 82 | peak / 1048576.0)); |
| 83 | } else { |
| 84 | memory = strings::StrCat(memory, "[", all.allocator_name(), |
| 85 | strings::Printf(" %.1fMB] ", tot / 1048576.0)); |
| 86 | } |
| 87 | } |
| 88 | } |
| 89 | const AttrSlice attrs(*node_); |
| 90 | string text; |
| 91 | if (IsSend(node_)) { |
| 92 | string tensor_name; |
| 93 | TF_CHECK_OK(GetNodeAttr(attrs, "tensor_name", &tensor_name)); |
| 94 | string recv_device; |
| 95 | TF_CHECK_OK(GetNodeAttr(attrs, "recv_device", &recv_device)); |
| 96 | text = strings::StrCat(memory, node_->name(), " = ", node_->op(), "(", |
| 97 | tensor_name, " @", recv_device); |
| 98 | } else if (IsRecv(node_)) { |
| 99 | string tensor_name; |
| 100 | TF_CHECK_OK(GetNodeAttr(attrs, "tensor_name", &tensor_name)); |
| 101 | string send_device; |
| 102 | TF_CHECK_OK(GetNodeAttr(attrs, "send_device", &send_device)); |
| 103 | text = strings::StrCat(memory, node_->name(), " = ", node_->op(), "(", |
| 104 | tensor_name, " @", send_device); |
| 105 | } else { |
| 106 | text = strings::StrCat(memory, node_->name(), " = ", node_->op(), "(", |
| 107 | absl::StrJoin(node_->input(), ", "), ")"); |
| 108 | } |
| 109 | stats_->set_timeline_label(text); |
| 110 | step_stats_collector_->Save(device, this); |
| 111 | } |
| 112 | |
| 113 | void NodeExecStatsWrapper::RecordExecutorStarted() { |
| 114 | int64 now_nanos = Env::Default()->NowNanos(); |