class_docstr: AWS Trainium stats.
code: def sample(self) -> None:
      try:
          raw_stats = json.loads(self.raw_samples[-1])
          neuron_runtime_data = [
              entry["report"]
              for entry in raw_stats["neuron_runtime_data"]
              if self._is_matching_entry(entry)
          ][
              0
          ]  # there should be only one entry with the pid
  
          neuroncores_in_use = neuron_runtime_data[
              "neuroncore_counters"
          ]["neuroncores_in_use"]
          # per-core utilization stats:
          neuroncore_utilization = {
              int(k): v["neuroncore_utilization"]
              for k, v in neuroncores_in_use.items()
          }
          # memory usage
          neuron_runtime_used_bytes = neuron_runtime_data[
              "memory_used"
          ]["neuron_runtime_used_bytes"]
          # memory usage totals
          host_total_memory_usage = neuron_runtime_used_bytes["host"]
          neuron_device_total_memory_usage = neuron_runtime_used_bytes[
              "neuron_device"
          ]
          # memory usage breakdown
          usage_breakdown = neuron_runtime_used_bytes["usage_breakdown"]
          host_memory_usage = _HostMemoryUsage(
              **usage_breakdown["host"]
          )
          neuroncore_memory_usage = {
              int(k): _NeuronCoreMemoryUsage(**v)
              for k, v in usage_breakdown[
                  "neuroncore_memory_usage"
              ].items()
          }
  
          # executed with torchrun? only keep the local_rank stats
          local_rank = int(os.environ.get("LOCAL_RANK", -1337))
          if local_rank >= 0:
              neuroncore_utilization = {
                  local_rank: neuroncore_utilization[local_rank]
              }
              neuroncore_memory_usage = {
                  local_rank: neuroncore_memory_usage[local_rank]
              }
  
          stats: _Stats = _Stats(
              neuroncore_utilization=neuroncore_utilization,
              host_total_memory_usage=host_total_memory_usage,
              neuron_device_total_memory_usage=neuron_device_total_memory_usage,
              host_memory_usage=host_memory_usage,
              neuroncore_memory_usage=neuroncore_memory_usage,
          )
          self.samples.append(stats)
  
      except Exception as e:  # noqa
          pass
  
imports: import collections
  import dataclasses
  import json
  import logging
  import os
  import pathlib
  import shutil
  import subprocess
  import sys
  import tempfile
  import threading
  import time
  from collections import deque
  from typing import TYPE_CHECKING, Any, Dict, List, Optional, Tuple, Union
  from wandb.sdk.lib import telemetry
  from .aggregators import aggregate_mean
  from .asset_registry import asset_registry
  from .interfaces import Interface, Metric, MetricsMonitor
label_desc: Use image node_id 7 for calling the NeuronCoreStats obj's underlying member method code with example usage: obj.sample() without return types