Add Google Cloud ML Diagnostics integration for profiling

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}"
 

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):

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 <RUN_NAME> <LOCATION> <PROJECT>
+
+set -euo pipefail
+
+if [[ $# -ne 3 ]]; then
+  echo "Usage: $0 <RUN_NAME> <LOCATION> <PROJECT>" >&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 } ]
+      }
+    }'

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,

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",
+    ],
+)
+

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):

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):

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":

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