Function UpdateStats

tensorflow/core/common_runtime/eager/kernel_and_device.cc:198–224

namespace {
void UpdateStats(OpKernelContext* context,
                 StepStatsCollector* step_stats_collector,
                 NodeExecStats* stats) {
  for (const auto& allocator_pair : context->ConsumeWrappedAllocators()) {
    AllocatorMemoryUsed* memory = stats->add_memory();
    memory->set_allocator_name(allocator_pair.first->Name());
    auto sizes = allocator_pair.second->GetSizes();
    memory->set_total_bytes(std::get<0>(sizes));
    memory->set_peak_bytes(std::get<1>(sizes));
    memory->set_live_bytes(std::get<2>(sizes));

    absl::optional<AllocatorStats> allocator_stats =
        allocator_pair.first->GetStats();
    if (allocator_stats) {
      memory->set_allocator_bytes_in_use(allocator_stats->bytes_in_use);
    }
    allocator_pair.second->GetRecordsAndUnRef();
  }
  auto* ms = stats->mutable_memory_stats();
  ms->set_temp_memory_size(context->temp_memory_allocated());
  for (const auto& alloc_id : context->persistent_alloc_ids()) {
    ms->mutable_persistent_tensor_alloc_ids()->Add(alloc_id);
  }

  ms->set_persistent_memory_size(context->persistent_memory_allocated());
  step_stats_collector->Finalize();
}

// In certain contexts (e.g. TPU async executions), the CancellationManager is
// used to shut down the device in error scenarios (as opposed to using the