From dabc9850052efb3671fee0222f6afc1489223294 Mon Sep 17 00:00:00 2001 From: rapatchi Date: Wed, 2 Sep 2026 07:16:09 +0000 Subject: [PATCH 1/2] feat(mldiag): Add Google Cloud ML Diagnostics profiling (xprof) support Introduce integration with `google_cloud_mldiagnostics` to support automated profiler (xprof) tracing on Google Cloud Platform. * Add `ManagedMLDiagnostics` singleton (`axlearn/common/managed_mldiagnostics.py`) to initialize a Google Cloud ML Diagnostics deferred run using `AXLEARN_JOB_NAME` and manage starting/stopping xprof tracing sessions. * Configure GKE launcher and TPU environment (`axlearn/cloud/gcp/tpu.py`, `axlearn/cloud/gcp/jobset_utils.py`) to inject `AXLEARN_JOB_NAME` into container environments and label pods with `managed-mldiagnostics-gke: true` when enabled. * Add `--enable_ml_diagnostics_xprof` and `--ml_diagnostics_region` flags to `axlearn/common/launch_trainer.py`. * Route profiler tracing in `SpmdTrainer` and `SpmdEvaler` through `ManagedMLDiagnostics` when `enable_ml_diagnostics_xprof` is active. * Add `axlearn/cloud/gcp/scripts/get_mldiag_urls.sh` helper script to extract profiler session TensorBoard URLs and Cloud Console URLs as JSON. * Add unit test coverage in `managed_mldiagnostics_test.py`, `trainer_test.py`, `evaler_test.py`, `launch_trainer_test.py`, and `jobset_utils_test.py`. TAG=agy CONV=2c9c481e-b3c4-463b-9064-6d33e10bb9cf --- axlearn/cloud/gcp/jobset_utils.py | 3 + axlearn/cloud/gcp/jobset_utils_test.py | 20 +++ axlearn/cloud/gcp/scripts/get_mldiag_urls.sh | 52 ++++++++ axlearn/cloud/gcp/tpu.py | 1 + axlearn/common/BUILD | 22 +++ axlearn/common/evaler.py | 50 +++++-- axlearn/common/evaler_test.py | 57 ++++++++ axlearn/common/launch_trainer.py | 16 +++ axlearn/common/launch_trainer_test.py | 48 +++++++ axlearn/common/managed_mldiagnostics.py | 102 ++++++++++++++ axlearn/common/managed_mldiagnostics_test.py | 133 +++++++++++++++++++ axlearn/common/trainer.py | 45 +++++-- axlearn/common/trainer_test.py | 55 ++++++++ pyproject.toml | 1 + 14 files changed, 578 insertions(+), 27 deletions(-) create mode 100755 axlearn/cloud/gcp/scripts/get_mldiag_urls.sh create mode 100644 axlearn/common/managed_mldiagnostics.py create mode 100644 axlearn/common/managed_mldiagnostics_test.py diff --git a/axlearn/cloud/gcp/jobset_utils.py b/axlearn/cloud/gcp/jobset_utils.py index d1c0cf6e6..5656a376b 100644 --- a/axlearn/cloud/gcp/jobset_utils.py +++ b/axlearn/cloud/gcp/jobset_utils.py @@ -863,6 +863,8 @@ def _build_pod(self) -> Nested[Any]: system = USER_FACING_NAME_TO_SYSTEM_CHARACTERISTICS[self._tpu_type] annotations, labels, selector, volumes, tolerations = {}, {}, {}, [], [] annotations["axlearn/main-container"] = cfg.name + if cfg.command and "enable_ml_diagnostics_xprof=True" in cfg.command: + labels["managed-mldiagnostics-gke"] = "true" volumes.append(dict(name="shared-output", emptyDir={})) if cfg.gcsfuse_mount: @@ -1165,6 +1167,7 @@ def _build_main_container(self) -> Nested[Any]: # These are common across all GPUReplicatedJobs, used for connecting between replicas env_vars: dict[str, Nested[str]] = {} + env_vars["AXLEARN_JOB_NAME"] = cfg.name env_vars["DISTRIBUTED_COORDINATOR"] = f"{cfg.name}-{cfg.job_name}-0-0.{cfg.name}:8080" env_vars["NUM_PROCESSES"] = f"{cfg.accelerator.num_replicas}" diff --git a/axlearn/cloud/gcp/jobset_utils_test.py b/axlearn/cloud/gcp/jobset_utils_test.py index 5c767773a..f3abeb1bf 100644 --- a/axlearn/cloud/gcp/jobset_utils_test.py +++ b/axlearn/cloud/gcp/jobset_utils_test.py @@ -801,6 +801,26 @@ def test_ephemeral_disk_mount_dataclass(self): self.assertEqual(m.size_gb, 500) self.assertEqual(m.read_only, False) + def test_mldiagnostics_label(self): + cases = [ + ("python3 -m trainer --enable_ml_diagnostics_xprof=True", True), + ("python3 -m trainer --enable_ml_diagnostics_xprof=False", False), + ("python3 -m trainer", False), + ] + for command, expected in cases: + with self._job_config(bundler_cls=ArtifactRegistryBundler) as (cfg, bundler_cfg): + cfg.set( + command=command, + output_dir="FAKE", + ) + job = cfg.instantiate(bundler=bundler_cfg.instantiate()) + pod = job._build_pod() # pylint: disable=protected-access + labels = pod["metadata"]["labels"] + if expected: + self.assertEqual(labels.get("managed-mldiagnostics-gke"), "true") + else: + self.assertNotIn("managed-mldiagnostics-gke", labels) + class CompositeReplicatedJobTest(TestCase): def test_composite_replicated_job(self): diff --git a/axlearn/cloud/gcp/scripts/get_mldiag_urls.sh b/axlearn/cloud/gcp/scripts/get_mldiag_urls.sh new file mode 100755 index 000000000..b8f024d7f --- /dev/null +++ b/axlearn/cloud/gcp/scripts/get_mldiag_urls.sh @@ -0,0 +1,52 @@ +#!/usr/bin/env bash +# +# Extracts profiler session TensorBoard URLs and Cloud Console URL +# for a given ML run as JSON. +# +# Usage: +# ./axlearn/cloud/gcp/scripts/get_mldiag_urls.sh + +set -euo pipefail + +if [[ $# -ne 3 ]]; then + echo "Usage: $0 " >&2 + echo "Example: $0 rapatchi-fuji-golden-mldiag-run us-central1 rapatchiexperiment-b6dk7n" >&2 + exit 1 +fi + +RUN="$1" +LOCATION="$2" +PROJECT="$3" + +# Resolve display name to resource ID if needed +RUN_ID=$(gcloud alpha mldiagnostics machine-learning-run list \ + --location="${LOCATION}" \ + --project="${PROJECT}" \ + --filter="displayName:${RUN} OR name:${RUN} OR workloadDetails.gke.id:${RUN}" \ + --format="value(name)" \ + --limit=1 \ + --quiet) + +if [[ -z "${RUN_ID}" ]]; then + RUN_ID="${RUN}" +fi + +HEX_ID="${RUN_ID##*/}" +CONSOLE_URL="https://console.cloud.google.com/cluster-director/diagnostics/details/${LOCATION}/${HEX_ID}?project=${PROJECT}" + +# Fetch profiler sessions as JSON and format with jq +gcloud alpha mldiagnostics profiler-session list \ + --machine-learning-run="${RUN_ID}" \ + --location="${LOCATION}" \ + --project="${PROJECT}" \ + --filter="dashboardUri:*" \ + --format=json \ + --quiet | jq \ + --arg run "${RUN}" \ + --arg console "${CONSOLE_URL}" \ + '{ + ($run): { + console_url: $console, + profiler_sessions: [ .[] | { ((.name | split("/") | last)): .dashboardUri } ] + } + }' diff --git a/axlearn/cloud/gcp/tpu.py b/axlearn/cloud/gcp/tpu.py index 2ab78e7c6..095727cc3 100644 --- a/axlearn/cloud/gcp/tpu.py +++ b/axlearn/cloud/gcp/tpu.py @@ -14,6 +14,7 @@ def get_default_env(*, tpu_type: str, num_tpu_slices: int, job_name: str) -> dict[str, Any]: """Gets the default environment for TPU pods.""" return dict( + AXLEARN_JOB_NAME=job_name, # Use a large refresh to mitigate DNS timeout issues until tf>2.12 upgrade. GCS_RESOLVE_REFRESH_SECS=600, TPU_TYPE=tpu_type, diff --git a/axlearn/common/BUILD b/axlearn/common/BUILD index 78a5759f2..8a7deded9 100644 --- a/axlearn/common/BUILD +++ b/axlearn/common/BUILD @@ -706,6 +706,7 @@ py_library( ":input_base", ":metrics", ":module", + ":managed_mldiagnostics", ":summary_writer", ":utils", "@axlearn_pip//absl_py", @@ -1205,6 +1206,7 @@ py_library( ":file_system", ":input_base", ":learner", + ":managed_mldiagnostics", ":measurement", ":module", ":optimizer_base", @@ -2129,6 +2131,7 @@ py_library( ":config", ":evaler", ":file_system", + ":managed_mldiagnostics", ":inference_output", ":input_base", ":layers", @@ -3475,3 +3478,22 @@ py_test( "@axlearn_pip//numpy", ], ) + +py_library( + name = "managed_mldiagnostics", + srcs = ["managed_mldiagnostics.py"], + visibility = ["//visibility:public"], + deps = [ + "@axlearn_pip//absl_py", + ], +) + +py_test( + name = "managed_mldiagnostics_test", + srcs = ["managed_mldiagnostics_test.py"], + deps = [ + ":managed_mldiagnostics", + "@axlearn_pip//absl_py", + ], +) + diff --git a/axlearn/common/evaler.py b/axlearn/common/evaler.py index ee764a603..699f65a76 100644 --- a/axlearn/common/evaler.py +++ b/axlearn/common/evaler.py @@ -19,6 +19,10 @@ from axlearn.common import flax_struct, input_base, summary_writer, utils from axlearn.common.base_model import BaseModel +from axlearn.common.managed_mldiagnostics import ( + MLDiagnosticsConfig, + is_ml_diagnostics_xprof_enabled, +) from axlearn.common.config import ( REQUIRED, InstantiableConfig, @@ -582,6 +586,8 @@ class Config(Module.Config): metric_calculator: BaseMetricCalculator.Config = ModelSummaryAccumulator.default_config() # If not None, writes input batches and `metric_calculator` forward outputs. output_writer: Optional[BaseOutputWriter.Config] = None + # Configuration for ML Diagnostics. + ml_diagnostics: Optional[MLDiagnosticsConfig] = None def __init__( self, @@ -615,6 +621,15 @@ def __init__( self._trace_steps = set() self._eval_policy: EvalPolicy = cfg.eval_policy.instantiate() + self._enable_ml_diagnostics_xprof: bool = is_ml_diagnostics_xprof_enabled( + cfg.ml_diagnostics + ) + if self._enable_ml_diagnostics_xprof: + from axlearn.common.managed_mldiagnostics import ManagedMLDiagnostics + try: + ManagedMLDiagnostics(cfg.ml_diagnostics) + except Exception: # pylint: disable=broad-except + pass def eval_step( self, @@ -691,25 +706,32 @@ def eval_step( "output was None at the end of a trace, not expected." ) jax.tree.map(lambda x: x.block_until_ready(), forward_outputs) - jax.profiler.stop_trace() + if self._enable_ml_diagnostics_xprof: + from axlearn.common.managed_mldiagnostics import ManagedMLDiagnostics + ManagedMLDiagnostics().stop_xprof() + else: + jax.profiler.stop_trace() self.vlog(2, "Stopped profiler tracing for evaler %s.", cfg.name) stop_trace_iter = None self._trace_steps.add(step) if batch_ix in cfg.trace_at_iters and len(self._trace_steps) <= 3: - try: - jax.profiler.start_trace(self.summary_writer.config.dir) - except RuntimeError as e: - if "Only one profile may be run at a time." in str(e): - # https://github.com/google/jax/blob/260f1d8b/jax/_src/profiler.py#L110-L111 - # No functionality is currently exposed to check this robustly. - raise RuntimeError( - "Nesting evaler profiling within a higher " - "level profile session is not currently supported. " - ) from e - # Else profiler is already running. - finally: - stop_trace_iter = batch_ix + 1 # We only look at one batch. + if self._enable_ml_diagnostics_xprof: + from axlearn.common.managed_mldiagnostics import ManagedMLDiagnostics + ManagedMLDiagnostics().start_xprof() + else: + try: + jax.profiler.start_trace(self.summary_writer.config.dir) + except RuntimeError as e: + if "Only one profile may be run at a time." in str(e): + # https://github.com/google/jax/blob/260f1d8b/jax/_src/profiler.py#L110-L111 + # No functionality is currently exposed to check this robustly. + raise RuntimeError( + "Nesting evaler profiling within a higher " + "level profile session is not currently supported. " + ) from e + # Else profiler is already running. + stop_trace_iter = batch_ix + 1 # We only look at one batch. self.vlog(2, "Start profiling evaler %s", cfg.name) with jax.profiler.StepTraceAnnotation(cfg.name, step_num=step): diff --git a/axlearn/common/evaler_test.py b/axlearn/common/evaler_test.py index 0e6121988..0bc822d2d 100644 --- a/axlearn/common/evaler_test.py +++ b/axlearn/common/evaler_test.py @@ -493,6 +493,63 @@ def fn(*, step, train_summaries) -> bool: ) self.assertIsNotNone(summaries) + def test_ml_diagnostics_xprof_tracing(self): + from unittest import mock + from axlearn.common.managed_mldiagnostics import ManagedMLDiagnostics, MLDiagnosticsConfig + ManagedMLDiagnostics._instance = None + + mesh = jax.sharding.Mesh(mesh_utils.create_device_mesh((1, 1)), ("data", "model")) + with mesh: + # Create model state. + model_cfg = DummyModel.default_config() + model = model_cfg.instantiate(parent=None) + model_param_partition_specs = jax.tree.map( + lambda spec: spec.mesh_axes, model.create_parameter_specs_recursively() + ) + model_state = pjit( + model.initialize_parameters_recursively, + in_shardings=(None,), + out_shardings=model_param_partition_specs, + )(jax.random.PRNGKey(0)) + + ManagedMLDiagnostics._instance = None + + mock_xprof_module = mock.MagicMock() + mock_xprof_class = mock_xprof_module.xprof + mock_xprof_inst = mock_xprof_class.return_value + + with mock.patch.dict("sys.modules", {"google_cloud_mldiagnostics": mock_xprof_module}), mock.patch.dict(os.environ, {"AXLEARN_JOB_NAME": "test_run"}): + with tempfile.TemporaryDirectory() as temp_dir: + evaler = ( + SpmdEvaler.default_config() + .set( + input=DummyInput.default_config().set( + total_num_batches=3, batch_size=4 + ), + name="spmd_evaler", + summary_writer=SummaryWriter.default_config().set(dir=temp_dir), + metric_calculator=DummyMetricCalculator.default_config(), + ml_diagnostics=MLDiagnosticsConfig( + enable_xprof=True, + region="us-central1", + gcs_path="gs://test/profiles", + ), + trace_at_iters=[1], + ) + .instantiate( + parent=None, + model=model, + model_param_partition_specs=model_param_partition_specs, + ) + ) + self.assertTrue(evaler._enable_ml_diagnostics_xprof) + + prng_key = jax.device_put(jax.random.PRNGKey(789), jax.NamedSharding(mesh, P())) + evaler.eval_step(1, prng_key=prng_key, model_params=model_state, force_run=True) + + mock_xprof_inst.start.assert_called_once() + mock_xprof_inst.stop.assert_called_once() + class ModelSummaryAccumulatorTest(absltest.TestCase): def test_accumulated_summaries_match(self): diff --git a/axlearn/common/launch_trainer.py b/axlearn/common/launch_trainer.py index 2727bc4b2..b46b4b313 100644 --- a/axlearn/common/launch_trainer.py +++ b/axlearn/common/launch_trainer.py @@ -12,6 +12,7 @@ from axlearn.common import file_system as fs from axlearn.common import measurement from axlearn.common.config import TrainerConfigFn, get_named_trainer_config +from axlearn.common.managed_mldiagnostics import MLDiagnosticsConfig from axlearn.common.trainer import SpmdTrainer, select_mesh_config from axlearn.common.utils import MeshShape, get_data_dir, infer_mesh_shape @@ -113,6 +114,16 @@ None, "The mesh selector string. See `SpmdTrainer.Config.mesh_rules` for details.", ) +flags.DEFINE_bool( + "enable_ml_diagnostics_xprof", + False, + "Whether to enable Google Cloud ML Diagnostics automated profiler (xprof) capture.", +) +flags.DEFINE_string( + "ml_diagnostics_region", + None, + "Google Cloud region for ML Diagnostics.", +) FLAGS = flags.FLAGS @@ -170,6 +181,11 @@ def get_trainer_config( ) if trainer_config.log_every_n_steps is None: trainer_config.log_every_n_steps = flag_values.trainer_log_every_n_steps + trainer_config.ml_diagnostics = MLDiagnosticsConfig( + enable_xprof=flag_values.enable_ml_diagnostics_xprof, + region=flag_values.ml_diagnostics_region, + gcs_path=f"{trainer_config.dir}/profiles", + ) for eval_cfg in trainer_config.evalers.values(): eval_cfg.trace_at_iters = [int(el) for el in flag_values.eval_trace_at_iters] if flag_values.device_monitor == "tpu": diff --git a/axlearn/common/launch_trainer_test.py b/axlearn/common/launch_trainer_test.py index c49282862..eac2aa2c7 100644 --- a/axlearn/common/launch_trainer_test.py +++ b/axlearn/common/launch_trainer_test.py @@ -200,6 +200,54 @@ def trainer_config_fn(): # Verify that the pre-existing value was not overridden self.assertEqual(cfg.crash_on_hang_timeout_seconds, 5000) + def test_ml_diagnostics_config(self): + """Test that ml_diagnostics config is properly set when enable_ml_diagnostics_xprof is True.""" + from axlearn.common.summary_writer import SummaryWriter + from axlearn.common.evaler import SpmdEvaler + + mock_trainer_config = mock.MagicMock() + mock_trainer_config.crash_on_hang_timeout_seconds = None + mock_trainer_config.watchdog_timeout_seconds = None + mock_trainer_config.dir = "/tmp/trainer" + mock_trainer_config.mesh_axis_names = None + mock_trainer_config.mesh_shape = None + mock_trainer_config.summary_writer = SummaryWriter.default_config() + + mock_evaler = SpmdEvaler.default_config() + mock_trainer_config.evalers = {"eval_test": mock_evaler} + mock_trainer_config.checkpointer = mock.MagicMock() + mock_trainer_config.checkpointer.trainer_dir = None + flag_values = { + "config": "config", + "config_module": "local_module", + "trainer_dir": "/tmp/trainer", + "enable_ml_diagnostics_xprof": True, + "ml_diagnostics_region": "us-central1", + "trainer_crash_on_hang_timeout_seconds": 9000, + "trainer_watchdog_timeout_seconds": 3600, + "trace_at_steps": [], + "eval_trace_at_iters": [], + "device_monitor": "none", + "mesh_selector": None, + } + fv = _flag_values_from_dict(flag_values) + + def trainer_config_fn(): + return mock_trainer_config + + with mock.patch( + f"{launch_trainer.__name__}.get_named_trainer_config", + side_effect=_mock_get_named_trainer_config, + ): + cfg = launch_trainer.get_trainer_config( + flag_values=fv, trainer_config_fn=trainer_config_fn + ) + + # Verify that ml_diagnostics is configured on the trainer itself (for profiling) + self.assertTrue(cfg.ml_diagnostics.enable_xprof) + self.assertEqual(cfg.ml_diagnostics.region, "us-central1") + self.assertEqual(cfg.ml_diagnostics.gcs_path, "/tmp/trainer/profiles") + if __name__ == "__main__": import sys diff --git a/axlearn/common/managed_mldiagnostics.py b/axlearn/common/managed_mldiagnostics.py new file mode 100644 index 000000000..0cb2459a6 --- /dev/null +++ b/axlearn/common/managed_mldiagnostics.py @@ -0,0 +1,102 @@ +# Copyright © 2026 Apple Inc. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# https://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Managed ML Diagnostics wrapper""" + +import logging +import os +import threading +from typing import Optional + +from axlearn.common.config import ConfigBase, config_class + + +class ManagedMLDiagnostics: + """Singleton wrapper for Google Cloud ML Diagnostics.""" + + _instance = None + _lock = threading.Lock() + + def __new__(cls, *args, **kwargs): + with cls._lock: + if cls._instance is None: + cls._instance = super(ManagedMLDiagnostics, cls).__new__(cls) + return cls._instance + + def __init__(self, cfg: Optional["MLDiagnosticsConfig"] = None): + with self._lock: + if hasattr(self, "_initialized"): + return + self._initialized = True + self._is_enabled = False + self._xprof = None + run_name = os.environ.get("AXLEARN_JOB_NAME", None) + + if cfg is None or run_name is None or not cfg.gcs_path: + return + if not cfg.enable_xprof: + return + try: + from google_cloud_mldiagnostics import machinelearning_run + # Initialize ML Diagnostics run (deferred) + kwargs = { + "name": run_name, + "gcs_path": cfg.gcs_path, + "on_demand_xprof": cfg.enable_xprof, + } + if cfg.region: + kwargs["region"] = cfg.region + machinelearning_run(**kwargs) + self._is_enabled = True + logging.info(f"Successfully created Google Cloud ML Diagnostics run (deferred): {run_name}") + + except Exception as e: + logging.error(f"Failed to start ML Diagnostics run (deferred): {e}", exc_info=True) + + def start_xprof(self): + """Starts ML diagnostics xprof tracing if available.""" + if not self._is_enabled: + return + with self._lock: + if self._xprof is None: + try: + from google_cloud_mldiagnostics import xprof as mldiag_xprof + self._xprof = mldiag_xprof() if mldiag_xprof is not None else None + except ImportError: + logging.warning( + "google-cloud-mldiagnostics is enabled but xprof import failed." + ) + return + if self._xprof is not None: + self._xprof._ensure_initialized() + self._xprof.start() + + def stop_xprof(self): + """Stops ML diagnostics xprof tracing.""" + if self._xprof is not None: + with self._lock: + self._xprof.stop() + + +@config_class +class MLDiagnosticsConfig(ConfigBase): + """Configuration for ML Diagnostics.""" + enable_xprof: bool = False + region: Optional[str] = None + gcs_path: Optional[str] = None + + +def is_ml_diagnostics_xprof_enabled(cfg: Optional[MLDiagnosticsConfig]) -> bool: + """Returns True if ML Diagnostics xprof profiling is configured and enabled.""" + return cfg is not None and cfg.enable_xprof diff --git a/axlearn/common/managed_mldiagnostics_test.py b/axlearn/common/managed_mldiagnostics_test.py new file mode 100644 index 000000000..4e8dec91f --- /dev/null +++ b/axlearn/common/managed_mldiagnostics_test.py @@ -0,0 +1,133 @@ +# Copyright © 2026 Apple Inc. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# https://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Tests for ManagedMLDiagnostics wrapper.""" + +import os +import sys +from unittest import mock +from absl import flags +from absl.testing import absltest, parameterized + +from axlearn.common.managed_mldiagnostics import ManagedMLDiagnostics, MLDiagnosticsConfig + + +class ManagedMLDiagnosticsTest(parameterized.TestCase): + + def setUp(self): + super().setUp() + ManagedMLDiagnostics._instance = None + + def test_initialize_run_success(self): + mock_machinelearning_run = mock.MagicMock() + modules_mock = { + "google_cloud_mldiagnostics": mock.MagicMock(), + "google_cloud_mldiagnostics.machinelearning_run": mock_machinelearning_run, + } + + with mock.patch.dict(sys.modules, modules_mock), mock.patch.dict(os.environ, {"AXLEARN_JOB_NAME": "test_run"}): + sys.modules["google_cloud_mldiagnostics"].machinelearning_run = mock_machinelearning_run + + cfg = MLDiagnosticsConfig( + gcs_path="gs://test", region="us-central1", enable_xprof=True + ) + diagnostics = ManagedMLDiagnostics(cfg) + mock_machinelearning_run.assert_called_once_with( + name="test_run", + region="us-central1", + gcs_path="gs://test", + on_demand_xprof=True, + ) + self.assertTrue(diagnostics._is_enabled) + + def test_initialize_run_only_once(self): + mock_machinelearning_run = mock.MagicMock() + modules_mock = { + "google_cloud_mldiagnostics": mock.MagicMock(), + "google_cloud_mldiagnostics.machinelearning_run": mock_machinelearning_run, + } + + with mock.patch.dict(sys.modules, modules_mock), mock.patch.dict(os.environ, {"AXLEARN_JOB_NAME": "test_run"}): + sys.modules["google_cloud_mldiagnostics"].machinelearning_run = mock_machinelearning_run + + cfg = MLDiagnosticsConfig(gcs_path="gs://test", region="us-central1", enable_xprof=True) + cfg2 = MLDiagnosticsConfig(gcs_path="gs://test2", region="us-central2", enable_xprof=True) + diagnostics = ManagedMLDiagnostics(cfg) + diagnostics2 = ManagedMLDiagnostics(cfg2) + + mock_machinelearning_run.assert_called_once_with( + name="test_run", + region="us-central1", + gcs_path="gs://test", + on_demand_xprof=True, + ) + + def test_initialize_run_import_error(self): + with mock.patch.dict(sys.modules, {"google_cloud_mldiagnostics": None}), mock.patch.dict(os.environ, {"AXLEARN_JOB_NAME": "test_run"}): + cfg = MLDiagnosticsConfig(gcs_path="gs://test", region="us-central1", enable_xprof=True) + diagnostics = ManagedMLDiagnostics(cfg) + self.assertFalse(diagnostics._is_enabled) + + def test_initialize_run_exception(self): + mock_machinelearning_run = mock.MagicMock(side_effect=RuntimeError("Some error")) + modules_mock = { + "google_cloud_mldiagnostics": mock.MagicMock(), + "google_cloud_mldiagnostics.machinelearning_run": mock_machinelearning_run, + } + + with mock.patch.dict(sys.modules, modules_mock), mock.patch.dict(os.environ, {"AXLEARN_JOB_NAME": "test_run"}): + sys.modules["google_cloud_mldiagnostics"].machinelearning_run = mock_machinelearning_run + + cfg = MLDiagnosticsConfig(gcs_path="gs://test", region="us-central1", enable_xprof=True) + diagnostics = ManagedMLDiagnostics(cfg) + self.assertFalse(diagnostics._is_enabled) + + def test_initialize_run_missing_name(self): + mock_machinelearning_run = mock.MagicMock() + modules_mock = { + "google_cloud_mldiagnostics": mock.MagicMock(), + "google_cloud_mldiagnostics.machinelearning_run": mock_machinelearning_run, + } + + with mock.patch.dict(sys.modules, modules_mock), mock.patch.dict(os.environ, {}, clear=True): + sys.modules["google_cloud_mldiagnostics"].machinelearning_run = mock_machinelearning_run + + cfg = MLDiagnosticsConfig(gcs_path="gs://test", region="us-central1", enable_xprof=True) + diagnostics = ManagedMLDiagnostics(cfg) + mock_machinelearning_run.assert_not_called() + self.assertFalse(diagnostics._is_enabled) + + def test_start_stop_xprof(self): + mock_xprof_cls = mock.MagicMock() + mock_xprof_inst = mock_xprof_cls.return_value + modules_mock = { + "google_cloud_mldiagnostics": mock.MagicMock(), + "google_cloud_mldiagnostics.xprof": mock_xprof_cls, + } + + with mock.patch.dict(sys.modules, modules_mock): + sys.modules["google_cloud_mldiagnostics"].xprof = mock_xprof_cls + diagnostics = ManagedMLDiagnostics() + diagnostics._is_enabled = True + + diagnostics.start_xprof() + mock_xprof_inst._ensure_initialized.assert_called_once() + mock_xprof_inst.start.assert_called_once() + + diagnostics.stop_xprof() + mock_xprof_inst.stop.assert_called_once() + + +if __name__ == "__main__": + absltest.main() diff --git a/axlearn/common/trainer.py b/axlearn/common/trainer.py index 632b3798d..dc3638234 100644 --- a/axlearn/common/trainer.py +++ b/axlearn/common/trainer.py @@ -51,6 +51,11 @@ from axlearn.common.optimizer_base import NestedOptParam, OptParam from axlearn.common.param_init import DefaultInitializer from axlearn.common.state_builder import Builder as TrainerStateBuilder +from axlearn.common.managed_mldiagnostics import ( + ManagedMLDiagnostics, + MLDiagnosticsConfig, + is_ml_diagnostics_xprof_enabled, +) from axlearn.common.summary_writer import BaseWriter, SummaryWriter from axlearn.common.update_transformation import ForwardOutputs # pytype: disable=pyi-error from axlearn.common.utils import ( @@ -236,6 +241,9 @@ class Config(Module.Config): # An optional recorder for measuring common metrics like step time. recorder: Optional[InstantiableConfig[measurement.Recorder]] = None + # Configuration for ML Diagnostics. + ml_diagnostics: Optional[MLDiagnosticsConfig] = None + # An additional context manager to run the training loop and initialization inside of. # The provided config should instantiate to a thunk that returns the context manager. context_manager: Optional[ConfigOr[Callable[[], ContextManager]]] = None @@ -328,6 +336,9 @@ def __init__( xsc_check_policy = maybe_instantiate(cfg.xsc_check_policy) self._xsc_check_policy: Optional[Callable[[int], bool]] = xsc_check_policy self._compiled_train_step: Optional[jax.stages.Compiled] = None + self._enable_ml_diagnostics_xprof: bool = is_ml_diagnostics_xprof_enabled( + cfg.ml_diagnostics + ) # Create all children within the mesh context so that utils.input_partition_spec() works # properly. @@ -381,6 +392,8 @@ def __init__( maybe_set_config( evaler_cfg.input, partition_spec=PartitionSpec(cfg.batch_axis_names) ) + if self._enable_ml_diagnostics_xprof: + evaler_cfg.ml_diagnostics = cfg.ml_diagnostics self._evalers[evaler_name] = self._add_child( evaler_name, evaler_cfg, @@ -1371,7 +1384,10 @@ def _maybe_stop_or_start_tracing( if self.step == stop_trace_step: assert output is not None jax.tree.map(lambda x: x.block_until_ready(), output) - jax.profiler.stop_trace() + if self._enable_ml_diagnostics_xprof: + ManagedMLDiagnostics().stop_xprof() + else: + jax.profiler.stop_trace() self._step_log("Stopped profiler tracing") updated_stop_trace_step = None @@ -1393,18 +1409,21 @@ def _maybe_stop_or_start_tracing( ) if should_start_tracing: self._step_log("Start profiler tracing") - profiler_options = jax.profiler.ProfileOptions() - if cfg.host_tracer_level is not None: - profiler_options.host_tracer_level = cfg.host_tracer_level - if cfg.device_tracer_level is not None: - profiler_options.device_tracer_level = cfg.device_tracer_level - if cfg.python_tracer_level is not None: - profiler_options.python_tracer_level = cfg.python_tracer_level - if cfg.tpu_trace_mode is not None: - profiler_options.advanced_configuration = {"tpu_trace_mode": cfg.tpu_trace_mode} - jax.profiler.start_trace( - self.summary_writer.config.dir, profiler_options=profiler_options - ) + if self._enable_ml_diagnostics_xprof: + ManagedMLDiagnostics().start_xprof() + else: + profiler_options = jax.profiler.ProfileOptions() + if cfg.host_tracer_level is not None: + profiler_options.host_tracer_level = cfg.host_tracer_level + if cfg.device_tracer_level is not None: + profiler_options.device_tracer_level = cfg.device_tracer_level + if cfg.python_tracer_level is not None: + profiler_options.python_tracer_level = cfg.python_tracer_level + if cfg.tpu_trace_mode is not None: + profiler_options.advanced_configuration = {"tpu_trace_mode": cfg.tpu_trace_mode} + jax.profiler.start_trace( + self.summary_writer.config.dir, profiler_options=profiler_options + ) updated_stop_trace_step = self.step + ( cfg.n_steps_for_each_trace if cfg.n_steps_for_each_trace is not None else 3 ) diff --git a/axlearn/common/trainer_test.py b/axlearn/common/trainer_test.py index 8ebd6be1b..70b8907e6 100644 --- a/axlearn/common/trainer_test.py +++ b/axlearn/common/trainer_test.py @@ -1387,6 +1387,61 @@ def test_ignore_undefined_loss(self, model_cfg: BaseModel.Config): trainer = cfg.instantiate(parent=None) trainer.run(prng_key=jax.random.PRNGKey(0)) + def test_ml_diagnostics_xprof_tracing(self): + from unittest import mock + from axlearn.common.managed_mldiagnostics import ManagedMLDiagnostics, MLDiagnosticsConfig + ManagedMLDiagnostics._instance = None + + cfg = self._trainer_config() + cfg.ml_diagnostics = MLDiagnosticsConfig( + enable_xprof=True, + region="us-central1", + gcs_path="gs://test/profiles", + ) + cfg.start_trace_steps = [2] + cfg.n_steps_for_each_trace = 2 + + mock_xprof_module = mock.MagicMock() + mock_xprof_class = mock_xprof_module.xprof + mock_xprof_inst = mock_xprof_class.return_value + + with mock.patch.dict("sys.modules", {"google_cloud_mldiagnostics": mock_xprof_module}), mock.patch.dict(os.environ, {"AXLEARN_JOB_NAME": "test_run"}): + ManagedMLDiagnostics(cfg.ml_diagnostics) + trainer = cfg.instantiate(parent=None) + self.assertTrue(trainer._enable_ml_diagnostics_xprof) + + # Initially, tracing is idle. + stop_trace_step = None + + # 1. At step 1: not in start_trace_steps. + trainer._step = 1 + stop_trace_step = trainer._maybe_stop_or_start_tracing(stop_trace_step, output=None) + self.assertIsNone(stop_trace_step) + mock_xprof_inst.start.assert_not_called() + + # 2. At step 2: start trace (since 2 is in start_trace_steps). + trainer._step = 2 + stop_trace_step = trainer._maybe_stop_or_start_tracing(stop_trace_step, output=None) + # stop_trace_step should be set to 2 + 2 = 4. + self.assertEqual(stop_trace_step, 4) + mock_xprof_inst.start.assert_called_once() + mock_xprof_inst.stop.assert_not_called() + + # 3. At step 3: tracing is active, but not at stop step. + trainer._step = 3 + stop_trace_step = trainer._maybe_stop_or_start_tracing(stop_trace_step, output=None) + self.assertEqual(stop_trace_step, 4) # remains unchanged. + mock_xprof_inst.stop.assert_not_called() + + # 4. At step 4: stop trace. + import jax.numpy as jnp + trainer._step = 4 + stop_trace_step = trainer._maybe_stop_or_start_tracing( + stop_trace_step, output={"loss": jnp.array(1.0)} + ) + self.assertIsNone(stop_trace_step) + mock_xprof_inst.stop.assert_called_once() + class SelectMeshConfigTest(test_utils.TestCase): def test_select_mesh_config(self): diff --git a/pyproject.toml b/pyproject.toml index 4bf8e03ee..e05cf1534 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -94,6 +94,7 @@ gcp = [ "pyOpenSSL>=22.1.0", # compat with cryptography version. "tpu-info==0.2.0", # For TPU monitoring from libtpu. https://github.com/AI-Hypercomputer/cloud-accelerator-diagnostics/tree/main/tpu_info "prometheus-client==0.21.0", # For TPU monitoring from tpu-device-plugin. + "google-cloud-mldiagnostics>=1.0.7", # For ML Diagnostics run logging. ] # For TPU training. # Note: Specify -f https://storage.googleapis.com/jax-releases/libtpu_releases.html during install. From 4a7bd30de1307a8cd8d0ea9a5f63af916c7b7440 Mon Sep 17 00:00:00 2001 From: rapatchi Date: Wed, 2 Sep 2026 07:25:35 +0000 Subject: [PATCH 2/2] feat(mldiag): Add Google Cloud ML Diagnostics metrics collection support Add support for recording and uploading training and evaluation metrics to Google Cloud ML Diagnostics. * Extend `ManagedMLDiagnostics` with a mapping dictionary `_METRIC_TO_METRIC_TYPE_NAME` and a `record_metric()` method to route AXLearn metrics (such as loss, learning rate, gradient norm, and step time) to GCP ML Diagnostics MetricTypes. * Add `MLDiagnosticsMetricsWriter` to `axlearn/common/summary_writer.py` and `inject_mldiagnostics_writer()` helper to automatically wrap existing summary writers into a `CompositeWriter`. * Inject `MLDiagnosticsMetricsWriter` into `SpmdTrainer` and `SpmdEvaler` when metrics recording is enabled. * Add `--enable_ml_diagnostics_metrics` flag to `axlearn/common/launch_trainer.py`. * Update GKE job launcher to recognize `enable_ml_diagnostics_metrics=True` when applying the `managed-mldiagnostics-gke` label. * Add unit test coverage in `managed_mldiagnostics_test.py`, `summary_writer_test.py`, `trainer_test.py`, `launch_trainer_test.py`, and `jobset_utils_test.py`. TAG=agy CONV=2c9c481e-b3c4-463b-9064-6d33e10bb9cf --- axlearn/cloud/gcp/jobset_utils.py | 2 +- axlearn/cloud/gcp/jobset_utils_test.py | 4 +- axlearn/common/evaler.py | 16 +++- axlearn/common/launch_trainer.py | 6 ++ axlearn/common/launch_trainer_test.py | 2 + axlearn/common/managed_mldiagnostics.py | 46 +++++++++- axlearn/common/managed_mldiagnostics_test.py | 47 +++++++++- axlearn/common/summary_writer.py | 89 ++++++++++++++++++ axlearn/common/summary_writer_test.py | 96 ++++++++++++++++++++ axlearn/common/trainer.py | 12 ++- axlearn/common/trainer_test.py | 65 +++++++++++++ 11 files changed, 374 insertions(+), 11 deletions(-) diff --git a/axlearn/cloud/gcp/jobset_utils.py b/axlearn/cloud/gcp/jobset_utils.py index 5656a376b..a5494bb2c 100644 --- a/axlearn/cloud/gcp/jobset_utils.py +++ b/axlearn/cloud/gcp/jobset_utils.py @@ -863,7 +863,7 @@ def _build_pod(self) -> Nested[Any]: system = USER_FACING_NAME_TO_SYSTEM_CHARACTERISTICS[self._tpu_type] annotations, labels, selector, volumes, tolerations = {}, {}, {}, [], [] annotations["axlearn/main-container"] = cfg.name - if cfg.command and "enable_ml_diagnostics_xprof=True" in cfg.command: + if cfg.command and ("enable_ml_diagnostics_metrics=True" in cfg.command or "enable_ml_diagnostics_xprof=True" in cfg.command): labels["managed-mldiagnostics-gke"] = "true" volumes.append(dict(name="shared-output", emptyDir={})) diff --git a/axlearn/cloud/gcp/jobset_utils_test.py b/axlearn/cloud/gcp/jobset_utils_test.py index f3abeb1bf..3d2bf3f68 100644 --- a/axlearn/cloud/gcp/jobset_utils_test.py +++ b/axlearn/cloud/gcp/jobset_utils_test.py @@ -803,8 +803,10 @@ def test_ephemeral_disk_mount_dataclass(self): def test_mldiagnostics_label(self): cases = [ + ("python3 -m trainer --enable_ml_diagnostics_metrics=True", True), ("python3 -m trainer --enable_ml_diagnostics_xprof=True", True), - ("python3 -m trainer --enable_ml_diagnostics_xprof=False", False), + ("python3 -m trainer --enable_ml_diagnostics_metrics=True --enable_ml_diagnostics_xprof=True", True), + ("python3 -m trainer --enable_ml_diagnostics_metrics=False", False), ("python3 -m trainer", False), ] for command, expected in cases: diff --git a/axlearn/common/evaler.py b/axlearn/common/evaler.py index 699f65a76..073beb64f 100644 --- a/axlearn/common/evaler.py +++ b/axlearn/common/evaler.py @@ -21,6 +21,7 @@ from axlearn.common.base_model import BaseModel from axlearn.common.managed_mldiagnostics import ( MLDiagnosticsConfig, + is_ml_diagnostics_metrics_enabled, is_ml_diagnostics_xprof_enabled, ) from axlearn.common.config import ( @@ -615,12 +616,11 @@ def __init__( model=model, model_param_partition_specs=model_param_partition_specs, ) - self._add_child("summary_writer", cfg.summary_writer) - if cfg.output_writer is not None: - self._add_child("output_writer", cfg.output_writer) - self._trace_steps = set() self._eval_policy: EvalPolicy = cfg.eval_policy.instantiate() + self._enable_ml_diagnostics_metrics: bool = is_ml_diagnostics_metrics_enabled( + cfg.ml_diagnostics + ) self._enable_ml_diagnostics_xprof: bool = is_ml_diagnostics_xprof_enabled( cfg.ml_diagnostics ) @@ -631,6 +631,14 @@ def __init__( except Exception: # pylint: disable=broad-except pass + writer_cfg = cfg.summary_writer + if self._enable_ml_diagnostics_metrics: + from axlearn.common.summary_writer import inject_mldiagnostics_writer + writer_cfg = inject_mldiagnostics_writer(writer_cfg, cfg.ml_diagnostics) + self._add_child("summary_writer", writer_cfg) + if cfg.output_writer is not None: + self._add_child("output_writer", cfg.output_writer) + def eval_step( self, step: int, diff --git a/axlearn/common/launch_trainer.py b/axlearn/common/launch_trainer.py index b46b4b313..4ba21c3fb 100644 --- a/axlearn/common/launch_trainer.py +++ b/axlearn/common/launch_trainer.py @@ -119,6 +119,11 @@ False, "Whether to enable Google Cloud ML Diagnostics automated profiler (xprof) capture.", ) +flags.DEFINE_bool( + "enable_ml_diagnostics_metrics", + False, + "Whether to enable Google Cloud ML Diagnostics metrics collection.", +) flags.DEFINE_string( "ml_diagnostics_region", None, @@ -183,6 +188,7 @@ def get_trainer_config( trainer_config.log_every_n_steps = flag_values.trainer_log_every_n_steps trainer_config.ml_diagnostics = MLDiagnosticsConfig( enable_xprof=flag_values.enable_ml_diagnostics_xprof, + enable_metrics=flag_values.enable_ml_diagnostics_metrics, region=flag_values.ml_diagnostics_region, gcs_path=f"{trainer_config.dir}/profiles", ) diff --git a/axlearn/common/launch_trainer_test.py b/axlearn/common/launch_trainer_test.py index eac2aa2c7..d5f233588 100644 --- a/axlearn/common/launch_trainer_test.py +++ b/axlearn/common/launch_trainer_test.py @@ -222,6 +222,7 @@ def test_ml_diagnostics_config(self): "config_module": "local_module", "trainer_dir": "/tmp/trainer", "enable_ml_diagnostics_xprof": True, + "enable_ml_diagnostics_metrics": True, "ml_diagnostics_region": "us-central1", "trainer_crash_on_hang_timeout_seconds": 9000, "trainer_watchdog_timeout_seconds": 3600, @@ -245,6 +246,7 @@ def trainer_config_fn(): # Verify that ml_diagnostics is configured on the trainer itself (for profiling) self.assertTrue(cfg.ml_diagnostics.enable_xprof) + self.assertTrue(cfg.ml_diagnostics.enable_metrics) self.assertEqual(cfg.ml_diagnostics.region, "us-central1") self.assertEqual(cfg.ml_diagnostics.gcs_path, "/tmp/trainer/profiles") diff --git a/axlearn/common/managed_mldiagnostics.py b/axlearn/common/managed_mldiagnostics.py index 0cb2459a6..694e4e7c8 100644 --- a/axlearn/common/managed_mldiagnostics.py +++ b/axlearn/common/managed_mldiagnostics.py @@ -17,11 +17,29 @@ import logging import os import threading -from typing import Optional +from typing import Any, Optional from axlearn.common.config import ConfigBase, config_class +# Mapping of concrete metric names to Google Cloud ML Diagnostics MetricType names. +_METRIC_TO_METRIC_TYPE_NAME = { + # 1. Model Quality + "loss": "LOSS", + "learning_rate": "LEARNING_RATE", + "learner/learning_rate": "LEARNING_RATE", + "learner/optimizer/learning_rate": "LEARNING_RATE", + "gradient_norm": "GRADIENT_NORM", + "learner/gradient_norm": "GRADIENT_NORM", + "learner/optimizer/gradient_norm": "GRADIENT_NORM", + "num_model_params": "TOTAL_WEIGHTS", + + # 2. Model Performance + "average_step_time": "STEP_TIME", + "step_time": "STEP_TIME", +} + + class ManagedMLDiagnostics: """Singleton wrapper for Google Cloud ML Diagnostics.""" @@ -45,7 +63,7 @@ def __init__(self, cfg: Optional["MLDiagnosticsConfig"] = None): if cfg is None or run_name is None or not cfg.gcs_path: return - if not cfg.enable_xprof: + if not (cfg.enable_xprof or cfg.enable_metrics): return try: from google_cloud_mldiagnostics import machinelearning_run @@ -64,6 +82,24 @@ def __init__(self, cfg: Optional["MLDiagnosticsConfig"] = None): except Exception as e: logging.error(f"Failed to start ML Diagnostics run (deferred): {e}", exc_info=True) + def record_metric(self, path: str, raw_value: Any, step: int): + """Record a training metric to ML Diagnostics.""" + if not self._is_enabled: + return + + try: + from google_cloud_mldiagnostics import metric_types, metrics as mldiag_metrics + path_lower = path.lower() + metric_val = float(raw_value) + + metric_type_name = _METRIC_TO_METRIC_TYPE_NAME.get(path_lower) + if metric_type_name: + metric_type = getattr(metric_types.MetricType, metric_type_name) + mldiag_metrics.record(metric_type, metric_val, step=step) + except Exception as e: + logging.error( + f"Failed to record metric {path} to ML Diagnostics: {e}", exc_info=True) + def start_xprof(self): """Starts ML diagnostics xprof tracing if available.""" if not self._is_enabled: @@ -93,6 +129,7 @@ def stop_xprof(self): class MLDiagnosticsConfig(ConfigBase): """Configuration for ML Diagnostics.""" enable_xprof: bool = False + enable_metrics: bool = False region: Optional[str] = None gcs_path: Optional[str] = None @@ -100,3 +137,8 @@ class MLDiagnosticsConfig(ConfigBase): def is_ml_diagnostics_xprof_enabled(cfg: Optional[MLDiagnosticsConfig]) -> bool: """Returns True if ML Diagnostics xprof profiling is configured and enabled.""" return cfg is not None and cfg.enable_xprof + + +def is_ml_diagnostics_metrics_enabled(cfg: Optional[MLDiagnosticsConfig]) -> bool: + """Returns True if ML Diagnostics metrics recording is configured and enabled.""" + return cfg is not None and cfg.enable_metrics diff --git a/axlearn/common/managed_mldiagnostics_test.py b/axlearn/common/managed_mldiagnostics_test.py index 4e8dec91f..68c8a235c 100644 --- a/axlearn/common/managed_mldiagnostics_test.py +++ b/axlearn/common/managed_mldiagnostics_test.py @@ -29,6 +29,51 @@ def setUp(self): super().setUp() ManagedMLDiagnostics._instance = None + @parameterized.parameters( + ("loss", 0.5, "LOSS"), + ("learning_rate", 0.1, "LEARNING_RATE"), + ("learner/learning_rate", 0.1, "LEARNING_RATE"), + ("learner/optimizer/learning_rate", 0.1, "LEARNING_RATE"), + ("gradient_norm", 1.2, "GRADIENT_NORM"), + ("learner/gradient_norm", 1.2, "GRADIENT_NORM"), + ("learner/optimizer/gradient_norm", 1.2, "GRADIENT_NORM"), + ("num_model_params", 1e9, "TOTAL_WEIGHTS"), + ("average_step_time", 0.5, "STEP_TIME"), + ("step_time", 0.5, "STEP_TIME"), + ("unmapped_metric", 123.4, None), + ) + def test_record_metric(self, path, raw_value, expected_metric_type_name): + mock_metric_types = mock.MagicMock() + # Set up mock enum values from the actual mapping + from axlearn.common.managed_mldiagnostics import _METRIC_TO_METRIC_TYPE_NAME + for name in set(_METRIC_TO_METRIC_TYPE_NAME.values()): + setattr(mock_metric_types.MetricType, name, name) + + mock_metrics = mock.MagicMock() + + # Mock the import of google_cloud_mldiagnostics + modules_mock = { + "google_cloud_mldiagnostics": mock.MagicMock(), + "google_cloud_mldiagnostics.metric_types": mock_metric_types, + "google_cloud_mldiagnostics.metrics": mock_metrics, + } + + with mock.patch.dict(sys.modules, modules_mock): + # Resolve dependencies in sys.modules + sys.modules["google_cloud_mldiagnostics"].metric_types = mock_metric_types + sys.modules["google_cloud_mldiagnostics"].metrics = mock_metrics + + diagnostics = ManagedMLDiagnostics() + diagnostics._is_enabled = True + diagnostics.record_metric(path, raw_value, step=10) + + if expected_metric_type_name is not None: + mock_metrics.record.assert_called_once_with( + expected_metric_type_name, float(raw_value), step=10 + ) + else: + mock_metrics.record.assert_not_called() + def test_initialize_run_success(self): mock_machinelearning_run = mock.MagicMock() modules_mock = { @@ -40,7 +85,7 @@ def test_initialize_run_success(self): sys.modules["google_cloud_mldiagnostics"].machinelearning_run = mock_machinelearning_run cfg = MLDiagnosticsConfig( - gcs_path="gs://test", region="us-central1", enable_xprof=True + gcs_path="gs://test", region="us-central1", enable_xprof=True, enable_metrics=True ) diagnostics = ManagedMLDiagnostics(cfg) mock_machinelearning_run.assert_called_once_with( diff --git a/axlearn/common/summary_writer.py b/axlearn/common/summary_writer.py index e84265836..3724e4a83 100644 --- a/axlearn/common/summary_writer.py +++ b/axlearn/common/summary_writer.py @@ -17,6 +17,7 @@ from axlearn.common import file_system as fs from axlearn.common.config import REQUIRED, ConfigBase, Required, RequiredFieldValue, config_class +from axlearn.common.managed_mldiagnostics import MLDiagnosticsConfig from axlearn.common.module import Module from axlearn.common.summary import AudioSummary, ImageSummary, Summary from axlearn.common.utils import Nested, NestedTensor, Tensor, tree_paths @@ -668,3 +669,91 @@ def _average_step_time(step_times: Sequence[float], window: int) -> Optional[flo if n <= 0: return None return sum(list(step_times)[-n:]) / n + + +class MLDiagnosticsMetricsWriter(BaseWriter): + """Writer for Google Cloud ML Diagnostics.""" + + @config_class + class Config(BaseWriter.Config): + """Configures MLDiagnosticsMetricsWriter. + + Attributes: + ml_diagnostics: Configuration for GCP ML Diagnostics. + write_every_n_steps: Writes summary every N steps. + """ + + dir: Optional[str] = "" + ml_diagnostics: Optional[MLDiagnosticsConfig] = None + write_every_n_steps: int = 1 + + def __init__(self, cfg: BaseWriter.Config, *, parent: Optional[Module]): + super().__init__(cfg, parent=parent) + cfg: MLDiagnosticsMetricsWriter.Config = self.config + from axlearn.common.managed_mldiagnostics import ( + is_ml_diagnostics_metrics_enabled, + ManagedMLDiagnostics, + ) + + self._ml_diagnostics = None + if is_ml_diagnostics_metrics_enabled(cfg.ml_diagnostics): + try: + self._ml_diagnostics = ManagedMLDiagnostics(cfg.ml_diagnostics) + except Exception: # pylint: disable=broad-except + pass + + def log_config(self, config: ConfigBase, step: int = 0): + pass + + def log_average_step_time(self, step: int, step_times: Sequence[float]): + cfg = self.config + if step % cfg.write_every_n_steps != 0: + return + average = _average_step_time(step_times, cfg.write_every_n_steps) + if average is not None: + self(step, {"average_step_time": average}) + + def __call__(self, step: int, values: dict[str, Any]) -> None: + if self._ml_diagnostics is None: + return + cfg: MLDiagnosticsMetricsWriter.Config = self.config + if step % cfg.write_every_n_steps != 0: + return + + prepared, paths = _prepare_for_d2h(values) + + def write(path: str, raw_value, value: Any): + if raw_value is None: + return + + self.vlog(3, "MLDiagnosticsWriter %s: %s=%s", self.path(), path, raw_value) + + if _match_summary_type("Scalar", value=value, raw_value=raw_value): + try: + self._ml_diagnostics.record_metric(path, raw_value, step) + except Exception: # pylint: disable=broad-except + pass + + jax.tree.map(write, paths, prepared, values, is_leaf=_is_summary_leaf) + + +def inject_mldiagnostics_writer( + writer_cfg: BaseWriter.Config, + ml_diagnostics: MLDiagnosticsConfig, +) -> BaseWriter.Config: + """Injects MLDiagnosticsMetricsWriter into a writer config, wrapping it in a CompositeWriter if needed.""" + mldiag_writer = MLDiagnosticsMetricsWriter.default_config().set( + ml_diagnostics=ml_diagnostics + ) + if isinstance(writer_cfg, CompositeWriter.Config): + writer_cfg.writers["mldiag"] = mldiag_writer + return writer_cfg + else: + return CompositeWriter.default_config().set( + dir=writer_cfg.dir, + writers={ + "tb": writer_cfg, + "mldiag": mldiag_writer + } + ) + diff --git a/axlearn/common/summary_writer_test.py b/axlearn/common/summary_writer_test.py index 184a2ff08..c95758ab3 100644 --- a/axlearn/common/summary_writer_test.py +++ b/axlearn/common/summary_writer_test.py @@ -26,9 +26,11 @@ from axlearn.common.evaler_test import DummyModel from axlearn.common.metrics import WeightedSummary from axlearn.common.summary import AudioSummary, ImageSummary +from axlearn.common.managed_mldiagnostics import MLDiagnosticsConfig from axlearn.common.summary_writer import ( CheckpointerAction, CompositeWriter, + MLDiagnosticsMetricsWriter, SummaryWriter, WandBWriter, _average_step_time, @@ -150,6 +152,69 @@ def test_log_config(self): self.assertEqual(expected, summaries) +class MLDiagnosticsMetricsWriterTest(TestCase): + """Tests MLDiagnosticsMetricsWriter.""" + + def test_record_mldiag_metrics(self): + cfg = MLDiagnosticsMetricsWriter.default_config().set( + name="test_mldiag", + write_every_n_steps=2, + ml_diagnostics=MLDiagnosticsConfig( + enable_metrics=True, + region="us-central1", + gcs_path="gs://test/profiles", + ) + ) + + mock_diagnostics_cls = mock.MagicMock() + mock_diagnostics_inst = mock_diagnostics_cls.return_value + + with mock.patch("axlearn.common.managed_mldiagnostics.ManagedMLDiagnostics", mock_diagnostics_cls): + writer = cfg.instantiate(parent=None) + mock_diagnostics_cls.assert_called_once_with(cfg.ml_diagnostics) + + # Step 1: write_every_n_steps == 2, so it shouldn't write on step 1 + writer( + step=1, + values={ + "loss": WeightedSummary(mean=3, weight=16), + }, + ) + mock_diagnostics_inst.record_metric.assert_not_called() + + # Step 2: write_every_n_steps == 2, so it should write + writer( + step=2, + values={ + "loss": WeightedSummary(mean=3, weight=16), + "accuracy": WeightedSummary(mean=0.7, weight=16), + "learner": {"learning_rate": 0.1}, + }, + ) + + expected_calls = [ + mock.call("loss", 3.0, 2), + mock.call("accuracy", 0.7, 2), + mock.call("learner/learning_rate", 0.1, 2), + ] + actual_calls = mock_diagnostics_inst.record_metric.call_args_list + self.assertEqual(len(expected_calls), len(actual_calls)) + mock_diagnostics_inst.record_metric.assert_has_calls(expected_calls, any_order=True) + + def test_disabled(self): + cfg = MLDiagnosticsMetricsWriter.default_config().set( + name="test_mldiag", + ml_diagnostics=MLDiagnosticsConfig(enable_metrics=False) + ) + + mock_diagnostics_cls = mock.MagicMock() + with mock.patch("axlearn.common.managed_mldiagnostics.ManagedMLDiagnostics", mock_diagnostics_cls): + writer = cfg.instantiate(parent=None) + mock_diagnostics_cls.assert_not_called() + + writer(step=1, values={"loss": 1.0}) + + class CompositeWriterTest(TestCase): """Tests CompositeWriter.""" @@ -203,6 +268,37 @@ def test_multiple_summary_writers_checkpoint(self): for sub_writer in writer.writers: sub_writer.log_checkpoint.assert_called_once() + def test_inject_mldiagnostics_writer(self): + from axlearn.common.summary_writer import inject_mldiagnostics_writer + from axlearn.common.managed_mldiagnostics import MLDiagnosticsConfig + + ml_diagnostics = MLDiagnosticsConfig( + enable_metrics=True, + region="us-central1", + gcs_path="gs://test/profiles", + ) + + # Case 1: Simple SummaryWriter config + tb_writer_cfg = SummaryWriter.default_config().set(dir="/tmp/tb") + injected_cfg = inject_mldiagnostics_writer(tb_writer_cfg, ml_diagnostics) + + self.assertTrue(issubclass(injected_cfg.klass, CompositeWriter)) + self.assertIn("tb", injected_cfg.writers) + self.assertIn("mldiag", injected_cfg.writers) + self.assertEqual(injected_cfg.writers["tb"].dir, "/tmp/tb") + self.assertEqual(injected_cfg.writers["mldiag"].ml_diagnostics, ml_diagnostics) + + # Case 2: Already a CompositeWriter config + composite_writer_cfg = CompositeWriter.default_config().set( + writers={"existing_tb": tb_writer_cfg} + ) + injected_cfg2 = inject_mldiagnostics_writer(composite_writer_cfg, ml_diagnostics) + + self.assertTrue(issubclass(injected_cfg2.klass, CompositeWriter)) + self.assertIn("existing_tb", injected_cfg2.writers) + self.assertIn("mldiag", injected_cfg2.writers) + self.assertEqual(injected_cfg2.writers["mldiag"].ml_diagnostics, ml_diagnostics) + class WandBWriterTest(TestCase): """Tests WandBWriter.""" diff --git a/axlearn/common/trainer.py b/axlearn/common/trainer.py index dc3638234..b61318957 100644 --- a/axlearn/common/trainer.py +++ b/axlearn/common/trainer.py @@ -54,6 +54,7 @@ from axlearn.common.managed_mldiagnostics import ( ManagedMLDiagnostics, MLDiagnosticsConfig, + is_ml_diagnostics_metrics_enabled, is_ml_diagnostics_xprof_enabled, ) from axlearn.common.summary_writer import BaseWriter, SummaryWriter @@ -336,6 +337,9 @@ def __init__( xsc_check_policy = maybe_instantiate(cfg.xsc_check_policy) self._xsc_check_policy: Optional[Callable[[int], bool]] = xsc_check_policy self._compiled_train_step: Optional[jax.stages.Compiled] = None + self._enable_ml_diagnostics_metrics: bool = is_ml_diagnostics_metrics_enabled( + cfg.ml_diagnostics + ) self._enable_ml_diagnostics_xprof: bool = is_ml_diagnostics_xprof_enabled( cfg.ml_diagnostics ) @@ -355,7 +359,11 @@ def __init__( cfg.summary_writer.dir = cfg.summary_writer.dir or os.path.join( cfg.dir, "summaries", "train_train" ) - self._add_child("summary_writer", cfg.summary_writer) + writer_cfg = cfg.summary_writer + if self._enable_ml_diagnostics_metrics: + from axlearn.common.summary_writer import inject_mldiagnostics_writer + writer_cfg = inject_mldiagnostics_writer(writer_cfg, cfg.ml_diagnostics) + self._add_child("summary_writer", writer_cfg) self._add_child("model", cfg.model) self._add_child("learner", cfg.learner) cfg.checkpointer.dir = cfg.checkpointer.dir or os.path.join(cfg.dir, "checkpoints") @@ -392,7 +400,7 @@ def __init__( maybe_set_config( evaler_cfg.input, partition_spec=PartitionSpec(cfg.batch_axis_names) ) - if self._enable_ml_diagnostics_xprof: + if self._enable_ml_diagnostics_metrics or self._enable_ml_diagnostics_xprof: evaler_cfg.ml_diagnostics = cfg.ml_diagnostics self._evalers[evaler_name] = self._add_child( evaler_name, diff --git a/axlearn/common/trainer_test.py b/axlearn/common/trainer_test.py index 70b8907e6..46c079d98 100644 --- a/axlearn/common/trainer_test.py +++ b/axlearn/common/trainer_test.py @@ -1442,6 +1442,71 @@ def test_ml_diagnostics_xprof_tracing(self): self.assertIsNone(stop_trace_step) mock_xprof_inst.stop.assert_called_once() + def test_ml_diagnostics_writer_injection(self): + from unittest import mock + from axlearn.common.managed_mldiagnostics import MLDiagnosticsConfig + from axlearn.common.summary_writer import CompositeWriter, MLDiagnosticsMetricsWriter, SummaryWriter + + # Reset singleton to avoid leakage + from axlearn.common.managed_mldiagnostics import ManagedMLDiagnostics + ManagedMLDiagnostics._instance = None + + cfg = self._trainer_config() + cfg.ml_diagnostics = MLDiagnosticsConfig( + enable_metrics=True, + region="us-central1", + gcs_path="gs://test/profiles", + ) + + # 1. Test case where original summary writer is a simple SummaryWriter + self.assertTrue(issubclass(cfg.summary_writer.klass, SummaryWriter)) + self.assertTrue(issubclass(cfg.evalers["eval_dummy"].summary_writer.klass, SummaryWriter)) + + mock_xprof_module = mock.MagicMock() + with mock.patch.dict("sys.modules", {"google_cloud_mldiagnostics": mock_xprof_module}), mock.patch.dict(os.environ, {"AXLEARN_JOB_NAME": "test_run"}): + trainer = cfg.instantiate(parent=None) + + # Verify that trainer's summary writer is wrapped in CompositeWriter + self.assertIsInstance(trainer.summary_writer, CompositeWriter) + self.assertIn("tb", trainer.summary_writer.children) + self.assertIn("mldiag", trainer.summary_writer.children) + self.assertIsInstance(trainer.summary_writer.children["mldiag"], MLDiagnosticsMetricsWriter) + + # Verify that evaler's summary writer is wrapped in CompositeWriter + evaler = trainer._evalers["eval_dummy"] + self.assertIsInstance(evaler.summary_writer, CompositeWriter) + self.assertIn("tb", evaler.summary_writer.children) + self.assertIn("mldiag", evaler.summary_writer.children) + self.assertIsInstance(evaler.summary_writer.children["mldiag"], MLDiagnosticsMetricsWriter) + self.assertEqual(evaler.config.ml_diagnostics, cfg.ml_diagnostics) + + # 2. Test case where original summary writer is already a CompositeWriter + ManagedMLDiagnostics._instance = None + cfg = self._trainer_config() + cfg.ml_diagnostics = MLDiagnosticsConfig( + enable_metrics=True, + region="us-central1", + gcs_path="gs://test/profiles", + ) + cfg.summary_writer = CompositeWriter.default_config().set( + writers={"tb": SummaryWriter.default_config()} + ) + cfg.evalers["eval_dummy"].summary_writer = CompositeWriter.default_config().set( + writers={"tb": SummaryWriter.default_config()} + ) + + with mock.patch.dict("sys.modules", {"google_cloud_mldiagnostics": mock_xprof_module}), mock.patch.dict(os.environ, {"AXLEARN_JOB_NAME": "test_run"}): + trainer = cfg.instantiate(parent=None) + + self.assertIsInstance(trainer.summary_writer, CompositeWriter) + self.assertIn("tb", trainer.summary_writer.children) + self.assertIn("mldiag", trainer.summary_writer.children) + + evaler = trainer._evalers["eval_dummy"] + self.assertIsInstance(evaler.summary_writer, CompositeWriter) + self.assertIn("tb", evaler.summary_writer.children) + self.assertIn("mldiag", evaler.summary_writer.children) + class SelectMeshConfigTest(test_utils.TestCase): def test_select_mesh_config(self):