google-cloud-mldiagnostics

What it is and what it does

Google Cloud ML Diagnostics is a Python SDK for integrating diagnostics into ML workloads running on Google Cloud. It lets you track machine learning runs, collect workload metrics (model quality, performance, system metrics), manage configs, and capture performance profiles using XProf—all visualized in the Google Cloud console. The package depends on google-api-core, google-auth, google-cloud-logging, psutil, packaging, and pynvml to communicate with Google Cloud services and gather system/GPU telemetry.

The SDK is designed for ML engineers to embed into their training or inference code on Google Cloud TPUs and GPUs. It works with orchestrators like Google Kubernetes Engine and custom setups. Metrics and profiles attach to a Machine Learning Run object, which can have multiple profiling sessions triggered either programmatically or on-demand from the UI. The package requires Python 3.8 or later and specific Google Cloud IAM permissions to function.

Use it for:

• Track metrics and profiles for JAX training runs on Google Cloud TPUs to identify performance bottlenecks.
• Collect system metrics (CPU, memory, GPU utilization) alongside model performance during inference workloads.
• Manage and version workload configs (software, system, user-defined) for reproducible ML experiments.
• Trigger XProf profiling sessions programmatically from within your ML code to capture performance data at key training steps.
• Collaborate on ML run diagnostics using shareable links to profiles and run metadata in the Google Cloud console.

Worth the install?

AI-flagged interpretation of the facts on this page — verify before relying

Collects metrics, configs, and performance profiles from ML workloads running on Google Cloud TPUs and GPUs, integrating with Google Cloud Logging and XProf for centralized diagnostics.

Yes, if you run ML workloads on Google Cloud TPUs or GPUs and need integrated diagnostics. The low install friction, active maintenance, and permissive license make it a straightforward add to JAX-based projects. Requires upfront Google Cloud setup (API enablement, IAM roles, Log Analytics) and is tightly coupled to Google Cloud infrastructure—not suitable for on-premises or non-Google-Cloud deployments.

Install

google-cloud-mldiagnostics on PyPI

pip install google-cloud-mldiagnostics

uv add google-cloud-mldiagnostics

poetry add google-cloud-mldiagnostics

Installing google-cloud-mldiagnostics

Before you install

Low friction: pure Python wheel with 6 runtime dependencies (google-api-core, google-auth, google-cloud-logging, psutil, packaging, pynvml). Active maintenance as of 9 days ago.

License in practice

Apache-2.0 permissive license allows commercial use, modification, and distribution with minimal restrictions—suitable for most production environments.

Quickstart

pip install google-cloud-mldiagnostics

from google.cloud import mldiagnostics
# Integrate with your ML workload to track metrics and profiles

Requires Google Cloud project with Cluster Director API enabled, appropriate IAM roles (clusterdirector.editor, logging.logWriter, storage.objectUser), and Log Analytics enabled on Cloud Logging _Default bucket.

Verify before relying

• Whether the package works with frameworks beyond JAX (description says 'JAX on Google Cloud TPUs and GPUs today').
• Whether multi-host profiling requires additional setup beyond the SDK itself.
• Performance overhead of metrics collection on large-scale workloads.

Package facts

Evidence: google_cloud_mldiagnostics-1.0.7-py3-none-any.whl

Tags

More Artificial Intelligence packages