From 4e2ddab6cc57835e26cffbc1b57ef070c393d3ec Mon Sep 17 00:00:00 2001 From: rapatchi Date: Wed, 2 Sep 2026 07:16:09 +0000 Subject: [PATCH] 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 Xprof 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..83df6d73c --- /dev/null +++ b/axlearn/cloud/gcp/scripts/get_mldiag_urls.sh @@ -0,0 +1,52 @@ +#!/usr/bin/env bash +# +# Extracts profiler session Xprof 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.