Managed profiling with XProf

You can use managed XProf profiling when you use ML Diagnostics workload monitoring, the ML Diagnostics CLI, or the ML Diagnostics SDK. Managed XProf is an XProf instance with a scalable backend that the ML Diagnostics platform provisions and manages in your Google Cloud project, whereas self-hosted XProf requires you to manually deploy and maintain your own XProf instances to view your profiles.

XProf is an open-source profiling and performance analysis tool for machine learning (ML) workloads and is part of the OpenXLA ecosystem. Compared to self-hosted XProf, managed XProf provides the following features:

  • Zero setup: Requires no manual installation or configuration of XProf visualization servers or their dependencies.
  • Centralized multi-run profile session access: Loads profile sessions across different ML runs with a single managed XProf instance per project.
  • Automatic scaling and fast loading: Accelerates the loading of large profiles and scales backend resources automatically when multiple users view profiles at the same time.
  • Built-in collaboration and security: Generates shareable links for profile sessions while protecting profiling instances from security vulnerabilities.
  • Multiple profile capture modes: Supports multi-host capture, on-demand profiling, continuous profiling in snapshot mode, and automatic profile capture from workload monitoring.
  • No additional infrastructure cost: Provides managed XProf compute resources at no extra charge, unlike self-hosting an XProf instance.

Deployment paths

You can deploy and enable managed XProf by using the ML Diagnostics SDK, workload monitoring, or the ML Diagnostics CLI.

The ML Diagnostics platform deploys one managed XProf instance with a scalable backend per project. You can use this instance to view all profile sessions captured across all ML runs in the project. Capturing profiles also requires a Cloud Storage bucket associated with the ML run to store the captured profile data.

You can enable managed XProf and associate a Cloud Storage bucket with your ML run in any of the following ways:

  • ML Diagnostics SDK: When you deploy your first ML run using the SDK, the ML Diagnostics platform automatically enables and deploys a managed XProf instance in your project. You specify the Cloud Storage bucket for storing profiles by setting the gcs_path parameter when you define the ML run in the SDK.
  • Workload monitoring: Workload monitoring doesn't enable managed XProf by default. Instead, the first time automatic profile capture triggers a profile session in your project, the ML Diagnostics platform automatically provisions and enables a managed XProf instance. By default, ML Diagnostics also provisions a regional Cloud Storage bucket in your project.
  • ML Diagnostics CLI: You can update any ML run to enable the XProf tool, which deploys the managed XProf instance in your project, and associate a Cloud Storage bucket with the ML run.

Capture profiles

ML Diagnostics supports four methods for capturing XProf profiles: programmatic capture, on-demand profiling, continuous profiling in snapshot mode, and automatic profile capture from workload monitoring. These methods differ in how a capture is triggered and what time window they record.

The following table compares the profile capture methods:

Capture method Description Supported frameworks
Programmatic capture

Embeds profiling commands directly in your workload code to start and stop recording profile data at specified training steps or functions.

Use this method when you want reproducible profiles of known training steps or specific code sections.

 JAX, vLLM, Tunix, TorchTPU
On-demand profiling

Triggers a live profile capture from the Google Cloud console or the gcloud CLI over a duration that you specify while your workload is running.

Use this method when you want to inspect a running workload in real time without modifying step ranges or restarting the workload.

 JAX, vLLM, Tunix, TorchTPU
Continuous profiling in snapshot mode (Preview)

Continuously records a 2 GB profile of device traces in a circular buffer on your TPU hosts and saves the latest 2 GB buffer to your Cloud Storage bucket when you trigger a snapshot.

Use this method when you need to inspect workload activity leading up to an intermittent slowdown or anomaly.

 JAX, vLLM, Tunix
Automatic profile capture from workload monitoring (Preview)

Automatically triggers an XProf profile capture when a workload monitoring performance degradation analyzer reports a "bad" state. You can configure automatic profile capture to trigger either an on-demand capture or, if continuous profiling is enabled, a continuous profiling snapshot.

Use this method when you want ML Diagnostics to capture trace data at the moment performance degrades without waiting for you to manually trigger a capture.

 JAX, vLLM, Tunix, TorchTPU

Programmatic capture

With programmatic capture, you embed profiling commands directly in your workload code to explicitly specify when to start and stop recording profile data. Any profiles written to the Cloud Storage bucket specified by the gcs_path parameter for your ML run automatically appear as profile sessions in your project's managed XProf instance, which you can view in the Google Cloud console or by using the Google Cloud CLI.

You can configure programmatic capture in either of the following ways:

On-demand profiling

On-demand profiling lets you trigger a profile capture in real time while your workload is running. Triggering an on-demand capture records a profile over a duration that you specify. You can enable on-demand profiling either directly in your workload code or by using the ML Diagnostics SDK.

Prerequisites

Before you capture an on-demand profile, complete the following setup steps:

  1. Start an XProf server on port 9999 in your workload for it to respond to on-demand capture triggers:

    • Directly in JAX: Call the profiler.start_server method on port 9999:

      profiler.start_server(port=9999)

    • ML Diagnostics SDK: Set on_demand_xprof=True in machinelearning_run when you define your ML run.

  2. Configure your cluster orchestrator or compute environment:

Capture an on-demand profile

You can trigger an on-demand profile capture from the Google Cloud console or the gcloud CLI. When you trigger a capture, specify the following parameters:

  • Duration: The capture duration, in milliseconds or seconds. The maximum duration that you can capture depends on your workload characteristics and how much data fits within the 2 GB profile limit.
  • Targets: The target nodes where you want to capture the profile:
    • Google Cloud console: All nodes running the workload are populated automatically, and you can select all or a subset of nodes.
    • gcloud CLI: List the profiler targets for the workload, then trigger a capture for all or a subset of those targets.
  • Device, host, and Python tracer levels: The level of detail to capture for the device, host, and Python traces.

Continuous profiling in snapshot mode

When continuous profiling is enabled, the ML Diagnostics platform continuously records a 2 GB profile in a circular buffer on your TPU hosts. When you trigger a snapshot, ML Diagnostics saves the latest 2 GB profile from the buffer as a profile session in your Cloud Storage bucket. This snapshot lets you inspect workload activity leading up to an issue. Continuous profiling only supports device tracing; it doesn't collect host or Python traces.

Prerequisites

Before you start continuous profiling, make sure that your environment meets the following requirements:

  • Supported library versions:
    • libtpu version 0.0.43 or later.
    • jax version 0.11.0 or later.
  • Cluster orchestrator or compute environment setup:

Start the XProf server in your workload code by calling profiler.start_server(port=9999), or set on_demand_xprof=True in machinelearning_run when using the ML Diagnostics SDK.

Start continuous profiling

You can start continuous profiling directly in your workload code, by using the gcloud CLI, or in the Google Cloud console. Once started, continuous profiling records 2 GB profiles in a circular buffer on the specified TPU hosts.

When you start continuous profiling using the gcloud CLI or the Google Cloud console, specify the following parameters:

  • Targets: The target nodes to profile continuously:
    • Google Cloud console: All nodes running the workload are populated automatically, and you can select all or a subset of nodes.
    • gcloud CLI: List the profiler targets for the workload, then start continuous profiling for all or a subset of those targets.
  • Device tracer level: The level of device tracing to capture continuously. Host and Python tracing are disabled during continuous profiling.

You can start continuous profiling directly from your Python workload code, in the gcloud CLI, or in the Google Cloud console.

Python workload code

To start continuous profiling directly from your Python workload code, call the start_continuous_profiling function:

from xprof.api import continuous_profiling_snapshot
continuous_profiling_snapshot.start_continuous_profiling('localhost:9999', {})

For more information, see XProf JAX profiling.

gcloud CLI

To start continuous profiling from the gcloud CLI, list the active profiler targets and start continuous profiling across those nodes using the following commands:

PROJECT_ID="PROJECT_ID"
LOCATION="LOCATION"
RUN_ID="RUN_ID"

# 1. Discover active ProfilerTarget IDs on the ML run:
TARGETS=$(gcloud alpha mldiagnostics profiler-target list \
  --machine-learning-run="${RUN_ID}" \
  --location="${LOCATION}" \
  --project="${PROJECT_ID}" \
  --format="csv[no-heading](name.basename())" | paste -sd, -)

# 2. Start continuous profiling across the target nodes:
gcloud alpha mldiagnostics machine-learning-run start-continuous-profiling \
  --machine-learning-run="${RUN_ID}" \
  --location="${LOCATION}" \
  --project="${PROJECT_ID}" \
  --targets="${TARGETS}" \
  --device-tracer-level=enabled

Replace the following:

  • PROJECT_ID: your Google Cloud project ID.
  • LOCATION: the region where your ML run is located.
  • RUN_ID: the ID of your ML run.
Google Cloud console

To start continuous profiling in the Google Cloud console, complete the following steps:

  1. In the Google Cloud console, go to Cluster Director or GKE > AI/ML > Run Diagnostics, and select your ML run.
  2. Open the Profiles tab.
  3. In the top action bar, click Start continuous profiling.
  4. Select the nodes to profile and the device tracer level, and then click Start.

After continuous profiling starts, the Capturing continuous profiles banner displays the start time of the session as Since START_TIME, and the action bar displays Take snapshot and Stop continuous profiling.

Capture a snapshot

While continuous profiling is running, you can capture a snapshot to save the latest 2 GB profile from the circular buffer as a profile session in your Cloud Storage bucket. You can capture a snapshot by using the gcloud CLI or in the Google Cloud console.

When you capture a continuous profiling snapshot, specify the following parameters:

  • Targets: All or a subset of the target nodes where continuous profiling is running:
    • Google Cloud console: All nodes configured for continuous profiling are populated automatically, and you can select all or a subset of nodes.
    • gcloud CLI: Specify all or a subset of the profiler targets where continuous profiling is running.
  • Profiler session name: A name for the captured snapshot session.

You can capture a snapshot by using the gcloud CLI or in the Google Cloud console.

gcloud CLI

To capture a continuous profiling snapshot using the gcloud CLI, run the following command:

PROJECT_ID="PROJECT_ID"
LOCATION="LOCATION"
RUN_ID="RUN_ID"

gcloud alpha mldiagnostics profiler-session capture "SNAPSHOT_NAME" \
  --machine-learning-run="${RUN_ID}" \
  --location="${LOCATION}" \
  --project="${PROJECT_ID}" \
  --targets="TARGET_NODES" \
  --kind=snapshot

Replace the following:

  • PROJECT_ID: your Google Cloud project ID.
  • LOCATION: the region where your ML run is located.
  • RUN_ID: the ID of your ML run.
  • SNAPSHOT_NAME: a name for the captured profile session.
  • TARGET_NODES: a comma-separated list of target nodes to capture.
Google Cloud console

To capture a continuous profiling snapshot in the Google Cloud console, complete the following steps:

  1. In the Google Cloud console, go to Cluster Director or GKE > AI/ML > Run Diagnostics, select your ML run, and open the Profiles tab.
  2. In the top action bar, click Take snapshot.
  3. In the Capture continuous profiling snapshot pane, enter a profile name, select the target nodes, and click Submit.

The snapshot appears in the Profiles table. If the link doesn't appear immediately, click Refresh.

Automatic profile capture from workload monitoring

With automatic profile capture from workload monitoring, ML Diagnostics automatically captures an XProf profile when a performance degradation analyzer reports a "bad" state, rather than waiting for you to manually trigger an on-demand capture after an issue occurs. You can configure automatic profile capture to trigger either an on-demand capture or, if continuous profiling is enabled, a continuous profiling snapshot. This process lets you diagnose the root cause of an issue using a profile captured at the time of the event.

For setup instructions and prerequisites, see Automatic profile capture in the workload monitoring documentation.

Agentic AI profile analysis

Agentic AI profile analysis automatically analyzes XProf performance profiles captured for TPU workloads and generates an HTML analysis report from raw XProf trace data. Using Gemini models on Vertex AI, the analysis agent runs XProf profiling tools, summarizes performance bottlenecks, and recommends code or configuration changes. When the required APIs are enabled, ML Diagnostics generates an agentic AI profile analysis report for every profile captured in the platform.

Agentic AI profile analysis uses the same underlying tools and capabilities as the open-source XProf skills repository. Rather than using the skill file directly or installing the XProf CLI required by the open-source skill, ML Diagnostics imports the underlying XProf tool implementations so that the agent can call them directly. As a result, the analysis output in ML Diagnostics is equivalent to running the open-source XProf skills with the XProf CLI on local xplane files.

Prerequisites

Before you use agentic AI profile analysis, complete the following setup steps in your Google Cloud project.

Enable required APIs

Enable the Vertex AI API (aiplatform.googleapis.com) so that the analysis agent can run Gemini models within your project:

  • Google Cloud console: In the API Library, enable the Vertex AI API.

  • gcloud CLI: Run the following command:

    gcloud services enable aiplatform.googleapis.com

Grant IAM roles

Grant your user account one of the following Identity and Access Management (IAM) roles:

  • Cluster Director Viewer (roles/hypercomputecluster.viewer): Grants read-only access to Cluster Director and ML Diagnostics resources, including generated analysis reports in the Google Cloud console under Run Diagnostics > Profiles.
  • Custom role (hypercomputecluster.locations.analysis): Grants permission to view generated analysis reports in the Google Cloud console under Run Diagnostics > Profiles when added to a custom role.

How it works

When a programmatic, on-demand, continuous, or automatic profile capture completes, an isolated background job runs the following analysis workflow on the XProf profile:

  1. Broad data collection: Runs 15 XProf tools in parallel to collect a health snapshot that includes step time, hardware utilization, memory usage, roofline bounds, slowest operations, and compiler inefficiencies. This step normalizes timings per training step and separates healthy operations from bottlenecks.
  2. Focused AI investigation: Reviews the collected health snapshot. If the workload has high compute utilization, low idle and communication overhead, and low out-of-memory risk, the agent skips additional tool calls and generates an OPTIMAL report. If bottlenecks are detected, the agent forms hypotheses and runs up to four targeted follow-up checks using additional XProf tools to inspect surrounding operations in the computation graph, timeline delays across hosts, and custom kernels.
  3. Report generation: Combines the metrics from XProf with the AI agent's root-cause analysis and recommended fixes. The generated summary.html report is saved to your Cloud Storage bucket and displayed in the Google Cloud console under Run Diagnostics > Profiles.

This workflow keeps your profile data and analysis private within your project. Raw profiler traces and generated HTML reports are stored in your Cloud Storage bucket, and all AI model calls run through Vertex AI in your Google Cloud project.

XProf tools

The following table describes the XProf tools that the agent runs during analysis:

Category XProf tools Description
Workload overview and step timing get_xprof_overview, get_step_trace, get_kpi_metrics, get_device_information, get_hosts, get_utilization_viewer, get_smart_suggestions

These tools run during broad data collection, and the get_hosts tool also runs during focused AI investigation, to check the following:

  • Hardware type, number of hosts and chips, and peak hardware limits
  • Average step time and breakdown of compute, network communication, and host or data-loading idle time
  • Matrix multiply unit (MXU), TensorCore, and memory bandwidth utilization, plus built-in XProf recommendations
Roofline and slowest operations get_roofline_model, get_hlo_op_profile

These tools run during broad data collection and focused AI investigation to check the following:

  • Whether the workload or specific operations are bound by compute or memory bandwidth
  • Time spent by operation type, such as matrix multiplication, attention, and collective operations, and links to Python source code lines
Compiler inefficiency detectors detect_layout_mismatch_copies, detect_unfused_reshapes, detect_unnecessary_convert_reduce, list_hlo_modules

These tools run during broad data collection to check the following:

  • Extra data copies caused by tensor layout mismatches between operations
  • Unfused reshape, transpose, or copy operations that waste memory bandwidth
  • Unnecessary precision conversions, such as bf16 to f32, inside reductions
Memory and OOM risk get_memory_profile, get_peak_allocations

These tools run during broad data collection to check the following:

  • Peak high-bandwidth memory (HBM) usage on the device compared to total capacity and memory fragmentation
  • Whether the out-of-memory (OOM) risk is LOW, MEDIUM, or HIGH, and the largest tensors in memory
Computation graph explorer get_hlo_neighborhood

This tool runs during focused AI investigation to check the operations immediately before and after a slow operation in the XLA graph and identify what blocks compiler fusion or forces data movement.
Timeline and kernel timing aggregate_xplane_events, get_kernel_stats

These tools run during focused AI investigation to check the following:

  • Timeline statistics across hosts to distinguish data-loader stalls from network communication delays
  • Execution time and call counts for individual device kernels
Custom kernel profiler llo_load, llo_query

These tools run during focused AI investigation to check low-level execution traces for custom kernels, such as Pallas or Mosaic, that appear opaque to standard high-level profiling.

AI profile analysis report

When analysis completes for a captured profile, ML Diagnostics saves the generated summary.html report to your Cloud Storage bucket and displays it in the Google Cloud console under Run Diagnostics > Profiles. You can use this report to review the overall health of your workload, inspect the top performance bottlenecks, and identify the recommended code or configuration changes.

Each generated report includes the following sections:

  • Analysis summary: Displays an overall status of OPTIMAL or NEEDS_ATTENTION, a 3-to-5-sentence summary, and key performance metrics such as step time, memory usage, and Roofline efficiency.
  • Prioritized bottlenecks and actionable fixes: Lists findings prioritized from P0 for critical issues to P2 for minor issues. Each finding includes:
    • What's happening: An explanation of the root cause, including tensor shapes and layouts when relevant.
    • Evidence: Direct citations of the XProf metrics and operations where the bottleneck was measured.
    • Expected impact and fix: The amount of step time or memory affected, paired with recommended framework, sharding, or compiler changes.
  • Detailed XProf tool analysis: Provides supporting data and detailed breakdowns from the following analyses:
    • Roofline analysis
    • Memory analysis
    • High-Level Optimizer (HLO) operation profile analysis
    • Kernel execution statistics
    • Custom kernel and low-level optimizer analysis