Add GPU monitor

Author: chunyang-wen

axlearn/cloud/gcp/monitoring/tpu_client.py

@@ -37,6 +37,8 @@
 
 
 # Interface names for libtpu metrics.
+# Reference:
+# https://github.com/AI-Hypercomputer/cloud-accelerator-diagnostics/blob/7d2b2921fc9393a3dec7be5440e25132c217549b/tpu_info/tpu_info/metrics.py#L29
 class MetricName(enum.Enum):
     """Metric names defined in libtpu."""
 
@@ -131,7 +133,7 @@ def get_chip_metrics(
     def sorted_metric_response(
         metric_name: str,
     ) -> list[tpu_metrics.Metric]:
-        # Manually annotate type until GRPC supports annotations
+        # Manually annotate type until GRPC supports annotations.
         # See https://github.com/grpc/grpc/issues/29041
         resp: tpu_metrics.MetricResponse = client.GetRuntimeMetric(
             tpu_metrics.MetricRequest(metric_name=metric_name)
@@ -153,7 +155,7 @@ def sorted_metric_response(
                 metric_results[i].hbm_memory_usage_bytes = metric.gauge.as_int
         elif metric_name == MetricName.TENSORCORE_DUTY_CYCLE_PERCENT:
             for i, metric in enumerate(metric_result):
-                metric_results[i].tensorcore_duty_cycle_percent = metric.gauge.as_double
+                metric_results[i].device_duty_cycle_percent = metric.gauge.as_double
 
     return metric_results
 
@@ -227,9 +229,9 @@ def get_chip_metrics_v2(
                         elif family.name == MetricV2Name.HBM_MEMORY_USAGE_BYTES.value:
                             metric_results[i].hbm_memory_usage_bytes = metric[2]
                         elif family.name == MetricV2Name.TENSORCORE_DUTY_CYCLE_PERCENT.value:
-                            metric_results[i].tensorcore_duty_cycle_percent = metric[2]
+                            metric_results[i].device_duty_cycle_percent = metric[2]
                         elif family.name == MetricV2Name.TENSORCORE_UTILIZATION.value:
-                            metric_results[i].tensorcore_utilization = metric[2]
+                            metric_results[i].device_utilization = metric[2]
                         elif family.name == MetricV2Name.HBM_MEMORY_BANDWIDTH_UTILIZATION.value:
                             metric_results[i].hbm_memory_bandwidth_utilization = metric[2]
 

axlearn/cloud/gcp/monitoring/tpu_client_test.py

@@ -80,7 +80,9 @@ def _(self, val: float):
         """Create Gauge from float."""
         return tpu_metrics.Gauge(as_double=val)
 
-    def GetRuntimeMetric(self, request: tpu_metrics.MetricRequest, context):
+    def GetRuntimeMetric(
+        self, request: tpu_metrics.MetricRequest, context
+    ):  # pylint: disable=unused-argument
         """Get the metric from the fake libtpu server."""
         metric_name = tpu_client.MetricName(request.metric_name)
         resp = self._responses[metric_name]
@@ -100,7 +102,9 @@ def GetRuntimeMetric(self, request: tpu_metrics.MetricRequest, context):
             )
         )
 
-    def ListSupportedMetrics(self, request: tpu_metrics.ListSupportedMetricsRequest, context):
+    def ListSupportedMetrics(
+        self, request: tpu_metrics.ListSupportedMetricsRequest, context
+    ):  # pylint: disable=unused-argument
         """List the supported metrics from the fake libtpu server."""
         # The test supported metrics are based on V5P libtpu.
         supported_metrics = [
@@ -210,7 +214,7 @@ def test_metrics(self, chip_type: device.TpuChip, responses):
             expected_usage = [
                 tpu_client.Usage(
                     device_id=i,
-                    tensorcore_duty_cycle_percent=d,
+                    device_duty_cycle_percent=d,
                     hbm_memory_usage_bytes=m,
                     hbm_memory_total_bytes=t,
                 )
@@ -313,8 +317,8 @@ def test_all(self):
         expected_usage = [
             tpu_client.Usage(
                 device_id=i,
-                tensorcore_duty_cycle_percent=100.0,
-                tensorcore_utilization=1.0 * (1 + i),
+                device_duty_cycle_percent=100.0,
+                device_utilization=1.0 * (1 + i),
                 hbm_memory_total_bytes=int(1.02803439616e11),
                 hbm_memory_usage_bytes=int(6.5e10),
                 hbm_memory_bandwidth_utilization=30.0,

axlearn/cloud/gcp/monitoring/tpu_device_monitor.py

@@ -68,10 +68,7 @@ def collect_metrics(self) -> list[Usage]:
     def is_host_idle(self, usages: list[Usage]) -> bool:
         """Check if the TPU device on the host are idle."""
         for usage in usages:
-            if (
-                usage.hbm_memory_bandwidth_utilization <= 0.1
-                and usage.tensorcore_utilization <= 0.1
-            ):
+            if usage.hbm_memory_bandwidth_utilization <= 0.1 and usage.device_utilization <= 0.1:
                 logging.info("TPU device %d is idle.", usage.device_id)
                 return True
         return False

axlearn/cloud/gcp/monitoring/tpu_device_monitor_test.py

@@ -22,8 +22,8 @@ def test_tpu_client(self):
                 device_id=i,
                 hbm_memory_total_bytes=int(1.02803439616e11),
                 hbm_memory_usage_bytes=int(6.5e10),
-                tensorcore_duty_cycle_percent=100.0,
-                tensorcore_utilization=1.0 * (1 + i),
+                device_duty_cycle_percent=100.0,
+                device_utilization=1.0 * (1 + i),
                 hbm_memory_bandwidth_utilization=30.0,
             )
             for i in range(4)

axlearn/common/launch_trainer.py

@@ -57,7 +57,7 @@
 flags.DEFINE_enum(
     "device_monitor",
     "none",
-    ["none", "tpu"],
+    ["none", "tpu", "gpu"],
     "Whether to enable the device monitor. "
     "The device monitor collects the system metrics and logs them periodically. "
     "The device monitor also logs the idle status of the devices on the host, "
@@ -116,6 +116,11 @@ def get_trainer_config(
         from axlearn.cloud.gcp.monitoring.tpu_device_monitor import create_tpu_monitor
 
         trainer_config.device_monitor = create_tpu_monitor()
+    elif flag_values.device_monitor == "gpu":
+        # pylint: disable-next=wrong-import-position,import-outside-toplevel
+        from axlearn.common.monitoring.gpu_device_monitor import create_gpu_monitor
+
+        trainer_config.device_monitor = create_gpu_monitor()
     if hasattr(trainer_config.checkpointer, "trainer_dir"):
         # Set trainer_dir if not already set.
         if not isinstance(trainer_config.checkpointer.trainer_dir, str):

axlearn/common/monitoring/device_monitor.py

@@ -94,7 +94,7 @@ def _start_monitoring(self):
         if self.config.check_interval_in_sec > 0:
             self._monitor_stopping = threading.Event()
             self._monitor_thread = threading.Thread(
-                name="tpu_device_monitor",
+                name="device_monitor",
                 target=self._monitor_loop,
             )
             self._monitor_thread.start()
@@ -115,4 +115,4 @@ def _monitor_loop(self):
             self._idle = self._check_host_and_log_metrics()
             if self._monitor_stopping.wait(timeout=self.config.check_interval_in_sec):
                 break
-        logging.info("mointor loop exit.")
+        logging.info("monitor loop exit.")

axlearn/common/monitoring/device_monitor_test.py

@@ -34,7 +34,7 @@ def is_host_idle(self, usages: list[Usage]) -> bool:
         # Make sure the usages are empty.
         return (
             usages[0].hbm_memory_bandwidth_utilization <= 0.1
-            and usages[0].tensorcore_utilization <= 0.1
+            and usages[0].device_utilization <= 0.1
         )
 
 
@@ -46,8 +46,8 @@ def test_client(self):
         fake_usage = [
             Usage(
                 device_id=0,
-                tensorcore_duty_cycle_percent=100.0,
-                tensorcore_utilization=1.0,
+                device_duty_cycle_percent=100.0,
+                device_utilization=1.0,
                 hbm_memory_total_bytes=100,
                 hbm_memory_usage_bytes=50,
                 hbm_memory_bandwidth_utilization=30.0,
@@ -71,8 +71,8 @@ def test_client_idle(self):
         fake_usage = [
             Usage(
                 device_id=0,
-                tensorcore_duty_cycle_percent=0.0,
-                tensorcore_utilization=0.0,
+                device_duty_cycle_percent=0.0,
+                device_utilization=0.0,
                 hbm_memory_total_bytes=100,
                 hbm_memory_usage_bytes=50,
                 hbm_memory_bandwidth_utilization=0.0,

axlearn/common/monitoring/gpu_client.py

@@ -0,0 +1,119 @@
+# Copyright © 2025 Apple Inc.
+
+"""Client for fetching GPU metrics via NVML."""
+import atexit
+
+from absl import logging
+
+
+class NVMLMetrics:
+    """NVMLMetrics provides interfaces to fetch GPU utilization/memory metrics via NVML.
+
+    Calling `pynvml.nvmlInit` multiple times will lead to potential issues and it should only
+    be called once.
+
+    And when the operations are completed, `pynvml.nvmlShutdown` should be called. Currently it is
+    called using `atexit`.
+    """
+
+    nvml_initialized = False
+    nvml = None
+
+    @classmethod
+    def init_nvml(cls):
+        """It is not thread-safe. Please see the docstring of the class for more details.
+
+        Users should not call `init_nvml` multiple times.
+        """
+        # pylint: disable-next=import-error,import-outside-toplevel
+        import pynvml as nvml  # pytype: disable=import-error
+
+        cls.nvml = nvml
+        if not cls.nvml_initialized:
+            try:
+                nvml.nvmlInit()
+            except:
+                logging.exception("Failed to initialize NVML Library for GPU metrics monitoring.")
+                raise
+            else:
+                cls.nvml_initialized = True
+                atexit.register(nvml.nvmlShutdown)
+
+    @classmethod
+    def get_gpu_device_count(cls):
+        cls.init_nvml()
+
+        try:
+            return cls.nvml.nvmlDeviceGetCount()
+        except:
+            logging.exception("Failed to get GPU device count.")
+            raise
+
+    @classmethod
+    def get_gpu_device_utilization(cls, device_id: int) -> float:
+        cls.init_nvml()
+
+        # pylint: disable-next=import-error,import-outside-toplevel
+        from pynvml import NVMLError  # pytype: disable=import-error
+
+        try:
+            device_handle = cls.nvml.nvmlDeviceGetHandleByIndex(device_id)
+
+            # Get all the utilization samples in the device buffer.
+            # Typically this covers about 10-13 seconds of data.
+            # Reference: https://docs.nvidia.com/deploy/nvml-api/group__nvmlDeviceQueries.html
+            # Search for nvmlDeviceGetSamples.
+            samples = cls.nvml.nvmlDeviceGetSamples(
+                device_handle, cls.nvml.NVML_GPU_UTILIZATION_SAMPLES, 0
+            )
+            util_samples = [sample.sampleValue.uiVal for sample in samples[1]]
+            if not util_samples:
+                logging.warning("No samples returned from pynvml.")
+                return 0
+            average_utilization = sum(util_samples) / len(util_samples)
+            return average_utilization
+        except NVMLError as e:
+            logging.exception("Failed to get GPU utilization metrics for device %d.", device_id)
+            logging.exception(e)
+            raise
+
+    @classmethod
+    def get_gpu_device_memory(cls, device_id: int) -> tuple[float, float]:
+        cls.init_nvml()
+
+        # pylint: disable-next=import-error,import-outside-toplevel
+        from pynvml import NVMLError  # pytype: disable=import-error
+
+        try:
+            device_handle = cls.nvml.nvmlDeviceGetHandleByIndex(device_id)
+            mem_info = cls.nvml.nvmlDeviceGetMemoryInfo(device_handle)
+
+            # Return tuple for memory usage, and total (in Bytes).
+            return mem_info.used, mem_info.total
+        except NVMLError as e:
+            logging.exception("Failed to get GPU memory info for device %d.", device_id)
+            logging.exception(e)
+            raise
+
+    @classmethod
+    def get_gpu_device_memory_utilization(cls, device_id: int) -> float:
+        cls.init_nvml()
+
+        # pylint: disable-next=import-error,import-outside-toplevel
+        from pynvml import NVMLError  # pytype: disable=import-error
+
+        try:
+            device_handle = cls.nvml.nvmlDeviceGetHandleByIndex(device_id)
+
+            # Get all the utilization samples in the device buffer.
+            # Typically this covers about 10-13 seconds of data.
+            samples = cls.nvml.nvmlDeviceGetSamples(
+                device_handle, cls.nvml.NVML_MEMORY_UTILIZATION_SAMPLES, 0
+            )
+            util_samples = [sample.sampleValue.uiVal for sample in samples[1]]
+            average_utilization = sum(util_samples) / len(util_samples)
+            return average_utilization
+        except NVMLError as e:
+            logging.exception("Failed to get GPU utilization metrics for device %d.", device_id)
+            logging.exception(e)
+            raise