| 210 | } |
| 211 | |
| 212 | Status SingleMachine::GetPeakMemoryUsage( |
| 213 | std::unordered_map<string, uint64>* device_peak_memory) const { |
| 214 | // Cpu_allocator->TracksAllocationSizes() returns true doesn't always mean the |
| 215 | // the AllocatorStats would be collected. |
| 216 | if (!cpu_allocator_stats_enabled_) { |
| 217 | return Status(error::INVALID_ARGUMENT, |
| 218 | "Tracking allocation for CPU is not enabled."); |
| 219 | } |
| 220 | |
| 221 | const DeviceMgr* device_mgr; |
| 222 | TF_RETURN_IF_ERROR(session_->LocalDeviceManager(&device_mgr)); |
| 223 | std::vector<Device*> devices = device_mgr->ListDevices(); |
| 224 | |
| 225 | device_peak_memory->clear(); |
| 226 | for (Device* device : devices) { |
| 227 | auto* allocator = device->GetAllocator(AllocatorAttributes()); |
| 228 | if (!allocator->TracksAllocationSizes()) { |
| 229 | return Status(error::INVALID_ARGUMENT, |
| 230 | "Tracking allocation is not enabled."); |
| 231 | } |
| 232 | absl::optional<AllocatorStats> stats = allocator->GetStats(); |
| 233 | (*device_peak_memory)[device->name()] = |
| 234 | (stats ? stats->peak_bytes_in_use : 0); |
| 235 | } |
| 236 | |
| 237 | return Status::OK(); |
| 238 | } |
| 239 | |
| 240 | Status SingleMachine::RunWithTimeout( |
| 241 | const std::vector<std::pair<string, Tensor>>& feed, |