ML Diagnostics platform

Note: Google Cloud ML Diagnostics is generally available for TPU workloads using JAX on Google Kubernetes Engine (GKE). Support for TPU workloads on vLLM, Tunix, and TorchTPU, and on Slurm and self-managed Kubernetes orchestrators, is in Preview.

Google Cloud ML Diagnostics is an end-to-end managed platform for optimizing, monitoring, and diagnosing AI and ML workloads on Google Cloud. Use ML Diagnostics to monitor workloads, as well as collect and visualize all workload metrics, configurations, and profiles within a single platform. ML Diagnostics is applicable to training, post-training reinforcement learning (RL), and inference workloads, and is compatible with all orchestrators on Cloud TPU. GKE support is generally available. Support for managed Slurm through Cluster Director, custom orchestrators such as Kubernetes and Slurm, and Compute Engine is in Preview.

ML Diagnostics includes the following generally available features:

• Workload monitoring: Automatically monitor and diagnose AI/ML workloads running on GKE TPUs. For each GKE job, ML Diagnostics Workload Monitoring automatically creates a Machine Learning run, monitors TPU duty cycle across the workload, detects events affecting the workload, and analyzes different layers of the workload -- such as infrastructure, orchestrator, and framework -- to help determine the potential causes of the event.
• Machine learning runs: Use ML Diagnostics to create and register your machine learning runs through Google Cloud CLI, or integrate the ML Diagnostics SDK with your workload. You can create and register runs automatically with workload monitoring, and deploy managed XProf instances with your machine learning runs, and collect and manage workload metrics, configurations, and profile sessions.
• gcloud CLI experience: Use the ML Diagnostics APIs through gcloud CLI to register and manage runs, deploy managed XProf resources, visualize profile sessions in storage buckets, and trigger profile captures from the CLI. You can also list all events detected by workload monitoring and obtain analyzer details for these events through the CLI or API.
• Python SDK: Use the open-source ML Diagnostics SDK integrated with ML workloads for a complete ML workload diagnostics experience. Monitor workloads and collect and manage workload metrics, configurations, and profiles on Google Cloud.
• Managed profiling: ML Diagnostics deploys a managed instance of XProf with a scalable backend into associated accounts, enabling the fast loading of large profiles. It supports multiple users simultaneously accessing profiles, and contains built-in features like multi-host profiling and on-demand profiling.
• Workload metrics: Track workload metrics, including model quality, model performance, and system metrics. With workload monitoring, you can get system metrics associated with the workload without integrating the SDK.
• Workload configuration management: Track workload configurations, including software configurations, system configurations, and user-defined configurations.
• Visualizations in Cluster Director and GKE: Visualize metrics, configurations, and profiles in Cluster Director and Google Kubernetes Engine in the Google Cloud console.
• Link sharing: Collaborate using shareable links with information on machine learning runs, workload monitoring, metrics, and profiles.

ML Diagnostics also includes the following Preview features:

• Workload monitoring: Automatically capture XProf profiles (with the Diagon DaemonSet) for further diagnosis when performance degradation analyzers show a bad state.
• Python SDK: Support for vLLM, Tunix, and TorchTPU workloads.
• Managed profiling: Continuous profiling in snapshot mode (with the Diagon DaemonSet) and agentic profile analysis.
• Orchestrators: Support for self-managed Kubernetes, Cluster Director-managed Slurm, self-managed Slurm, and Compute Engine.

Workload frameworks and orchestrator types

ML Diagnostics supports JAX and is generally available on GKE. Workload frameworks like TorchTPU, vLLM, and Tunix are available in Preview.

The following table summarizes the supported features and required components for TPU workloads across orchestrator types:

Orchestrator type	Supported features and required components
GKE	Workload monitoring for GKE Job and JobSet workloads SDK metrics and programmatic profiling (GKE webhook) On-demand profiling (GKE operator) On-demand profiling, continuous profiling, and automatic profile capture from workload monitoring with the Diagon DaemonSet (Preview)
Self-managed Kubernetes (Preview)	SDK metrics and programmatic profiling (Kubernetes webhook) On-demand profiling (Kubernetes operator) On-demand and continuous profiling with the Diagon DaemonSet (Preview)
Cluster Director-managed Slurm, self-managed Slurm, and Compute Engine (Preview)	SDK metrics and programmatic profiling On-demand and continuous profiling with the Diagon DaemonSet (Preview)

To request access to Preview features, contact your account team. You don't need to request access to use SDK metrics, programmatic profiling, or on-demand profiling with TorchTPU, vLLM, and Tunix workloads on GKE.

User paths

You can use the ML Diagnostics platform automatically with workload monitoring for all GKE jobs, or integrate the SDK with your workload for the complete ML workload diagnostics experience. You can interact with the ML Diagnostics platform through Google Cloud console or the CLI. Workload monitoring is automatically available for all workloads deployed with GKE on TPUs.

With the CLI, you can use ML Diagnostics gcloud CLI to create or modify a machine learning run, and deploy the managed XProf resources. With the ML Diagnostics SDK, the SDK needs to be integrated into your ML workload to collect and manage workload metrics and configurations, and deploy managed XProf resources.

To get started, use one of the following guides:

• Monitor workloads
• Get started with SDK
• Get started with CLI

Prerequisites

Before using ML Diagnostics, enable the Cluster Director API and add the required IAM permissions. If you are using GKE, workload monitoring is already enabled. However, the ML Diagnostics SDK and on-demand profiling require additional GKE artifacts and configuration of your GKE cluster. For more information, see Set up GKE.

Enable Cluster Director API

Note: The Cluster Director API is not related to TPU Cluster Director. Google Cloud ML Diagnostics relies on the Cluster Director API for accessing gcloud CLI commands; for collecting metrics, configurations, and profiles in the platform; and for accessing the Google Cloud console.

You don't need to use Cluster Director for deploying and managing your clusters in order to use the ML Diagnostics product. ML Diagnostics works with clusters managed by GKE, Cluster Director, or custom orchestrators. ML Diagnostics is part of the Cluster Director family of APIs, but doesn't depend on users using the Cluster Director product itself.

For more information about enabling Cluster Director API, see Enabling an API in your Google Cloud project.

IAM permissions

The Google Cloud service account used by your workload requires the following IAM roles assigned on the project:

• roles/hypercomputecluster.editor: For full access to create and manage MLRun resources and view the user interface.
• roles/logging.logWriter: To write logs and metrics to Cloud Logging.
• roles/storage.objectUser: To save profiles to the Cloud Storage bucket specified in machinelearning_run.

You can also create custom roles that assign users a subset of APIs needed for ML Diagnostics. Since users don't need to use Cluster Director for managing clusters, you only need to assign the following permissions:

• hypercomputecluster.locations.get
• hypercomputecluster.locations.list
• hypercomputecluster.clusters.list
• hypercomputecluster.nodes.list
• hypercomputecluster.machineLearningRuns.create
• hypercomputecluster.machineLearningRuns.delete
• hypercomputecluster.machineLearningRuns.get
• hypercomputecluster.machineLearningRuns.list
• hypercomputecluster.machineLearningRuns.update
• hypercomputecluster.operations.cancel
• hypercomputecluster.operations.delete
• hypercomputecluster.operations.get
• hypercomputecluster.operations.list
• resourcemanager.projects.get
• resourcemanager.projects.list

For workloads on Google Kubernetes Engine, use Workload Identity Federation to associate a Kubernetes Service Account with a Google Cloud service account that has been granted the required roles. Pods running on the host network (hostNetwork: true) don't use Workload Identity Federation for GKE. For these cases, see Alternatives to Workload Identity Federation for GKE.

Enable Observability Analytics on Cloud Logging

The model metrics and performance metrics time-series charts render their data using an Observability Analytics query against the project's _Default log bucket. If Observability Analytics isn't enabled on that bucket, these charts fail and the widget shows a Something went wrong in the backend or 0 time series message, even though the metric data is present in Cloud Logging and profiles and system metrics render as expected.

Enable Observability Analytics once per project. This upgrade is non-destructive, doesn't incur additional cost, and can't be disabled. You can still query existing logs:

gcloud logging buckets update _Default --location=global \
  --enable-analytics --project=PROJECT_ID

This command requires the logging.buckets.update permission, which is included in roles such as roles/logging.configWriter.

Note: Observability Analytics doesn't backfill data. Only logs ingested after you enable Observability Analytics are queryable. Runs that completed before you enabled it still show empty charts, so start a new run to verify.

Pricing

You are charged for the storage of metrics through Cloud Logging, and the storage of profiles through Cloud Storage. There is no need to enable any extra billing for these services when using the ML Diagnostics platform. There is no charge for the managed XProf resources deployed by the ML Diagnostics platform.