Monitor goodput with the ML Goodput Measurement library

The ML Goodput Measurement library (ml-goodput-measurement) is a Python package that helps you measure the efficiency of your ML training workloads running on Cloud TPU VMs. The library provides metrics that measure workload goodput, which is the proportion of TPU usage time spent making productive, preserved training progress. Conversely, badput is the proportion of total time spent on non-productive activities like startup overhead, I/O stalls, and disruption recovery.

You can visualize goodput metrics in realtime with Cloud Monitoring dashboards and TensorBoard, enabling you to pinpoint bottlenecks, optimize resource utilization, and ultimately reduce training costs.

For more information, see the ML Goodput Measurement GitHub repository.

Goodput metrics

The ML Goodput Measurement library provides the following metrics, which are also available to view in Cloud Monitoring and TensorBoard. The metrics in this table must be prefixed with compute.googleapis.com/workload/. For example, the full metric name for goodput_time is compute.googleapis.com/workload/goodput_time.

  • goodput_time: The total productive training time in seconds. This can be interpreted as the cumulative goodput.

  • badput_time: The total unproductive training time in seconds (startup, stalls, recovery). This can be interpreted as the cumulative badput.

  • total_elapsed_time: The total elapsed time (wall-clock duration) of the workload in seconds. The elapsed time is measured from the time the application starts to either the current time or the time of job completion.

  • interval_goodput: The goodput rate over a specified time period (for example, the last 24 hours). The metric provides a rolling window for goodput.

  • interval_badput: The badput rate over a specified time period. The metric provides a rolling window for badput, and is useful for identifying transient issues like spikes in I/O operations.

  • disruptions: The cumulative count of job disruptions. Disruptions are events that cause the training process to stop unexpectedly, requiring a restart. For example, hardware failures and maintenance events.

  • step_time_deviation: The amount of unproductive time due to variation in training step times. This is also known as "jitter". The metric measures spike-sensitive stability, which is the deviation of recent "worst-case" steps against a historical baseline in seconds. A step is a single iteration of the training loop.

  • performance: The estimated fastest stable step time (baseline) in seconds. The ideal step time is the fastest time a single training step can take under optimal conditions, free from transient noise or jitter.

  • max_productive_steps: The highest step count reached that was successfully preserved.

Installation

Use the following steps to set up the ML Goodput Measurement library with your TPU workloads:

  1. Enable the Cloud Logging API and Cloud Monitoring API.
  2. If deploying on Google Kubernetes Engine (GKE), configure all node pools with the cloud-platform access scope.

Install the ml-goodput-measurement package on your training host and analysis machine:

pip install ml-goodput-measurement

Monitor with the Goodput library

To use the ML Goodput Measurement library, initialize a GoodputRecorder instance, instrument your training code by wrapping key sections with the record_event context manager, and monitor metrics in realtime with GoodputMonitor. The monitor runs a background process to periodically calculate goodput metrics from the recorded events, and uploads them to Cloud Monitoring and TensorBoard for real-time analysis and visualization.

Initialize the Goodput recorder

Initialize the GoodputRecorder, which is the core component of the ML Goodput Measurement library.

import datetime
import jax