Fix multi-host profiling and resolve profiler tracer collisions

src/maxdiffusion/max_utils.py

@@ -85,7 +85,7 @@ def _jax_profiler_enabled(config):
 
 
 def _ml_diagnostics_profiler_enabled(config):
-  return "enable_ml_diagnostics" in config.get_keys() and config.enable_ml_diagnostics and jax.process_index() == 0
+  return "enable_ml_diagnostics" in config.get_keys() and config.enable_ml_diagnostics
 
 
 def profiler_enabled(config):
@@ -120,22 +120,27 @@ class Profiler:
   def __init__(self, config, session_name=None):
     self.config = config
     self.session_name = session_name
+    self.mld_xprof = None
+    self._active = None  # "mld" | "jax" | None
 
   def start(self):
-    if _jax_profiler_enabled(self.config) and _ml_diagnostics_profiler_enabled(self.config):
+    use_mld = _ml_diagnostics_profiler_enabled(self.config) and xprof is not None
+    use_jax = _jax_profiler_enabled(self.config)
+
+    if use_mld and use_jax:
       max_logging.log(
-          "Warning: Both ML Diagnostics profiler and JAX profiler are enabled. "
-          "This may cause increased overhead and duplicate profiling data. It "
-          "is recommended to enable only one profiler at a time for accurate "
-          "performance analysis."
+          "Both ML Diagnostics and JAX profilers are enabled. They share the same "
+          "underlying JAX tracer and cannot run concurrently. Using ML Diagnostics "
+          "and skipping the standalone JAX profiler."
       )
+      use_jax = False
 
-    if _ml_diagnostics_profiler_enabled(self.config) and xprof is not None:
+    if use_mld:
       ensure_machinelearning_job_runs(self.config)
       self.mld_xprof = xprof()
       self.mld_xprof.start(self.session_name)
-
-    if _jax_profiler_enabled(self.config):
+      self._active = "mld"
+    elif use_jax:
       log_dir = self.config.tensorboard_dir
       if log_dir.startswith("gs://"):
         log_dir = os.path.join("/tmp/profiler_traces", self.config.run_name)
@@ -144,14 +149,13 @@ def start(self):
       os.makedirs(log_dir, exist_ok=True)
       max_logging.log(f"Starting profiler trace in: {log_dir}")
       jax.profiler.start_trace(log_dir)
+      self._active = "jax"
 
   def stop(self):
-    if _ml_diagnostics_profiler_enabled(self.config) and xprof is not None:
-      ensure_machinelearning_job_runs(self.config)
+    if self._active == "mld":
       if self.mld_xprof is not None:
         self.mld_xprof.stop()
-
-    if _jax_profiler_enabled(self.config):
+    elif self._active == "jax":
       jax.profiler.stop_trace()
 
       trace_dir = self.config.tensorboard_dir
@@ -171,6 +175,7 @@ def stop(self):
               blob = bucket.blob(blob_name)
               blob.upload_from_filename(local_file)
               max_logging.log(f"Uploaded {local_file} to gs://{bucket_name}/{blob_name}")
+    self._active = None
 
   def __enter__(self):
     self.start()

src/maxdiffusion/tests/profiler_test.py

@@ -63,6 +63,52 @@ def test_ml_diagnostics_profiler(self, mock_process_index, mock_xprof, mock_ml_r
 
     mock_xprof.return_value.stop.assert_called_once()
 
+  @patch("maxdiffusion.max_utils.machinelearning_run")
+  @patch("maxdiffusion.max_utils.xprof")
+  @patch("jax.process_index", return_value=1)
+  def test_ml_diagnostics_profiler_non_master_host(
+      self, mock_process_index, mock_xprof, mock_ml_run
+  ):
+    """Tests that ML Diagnostics profiler is also enabled on non-master hosts (process_index != 0)."""
+    config = MockConfig(
+        enable_ml_diagnostics=True,
+        profiler_gcs_path="gs://fake-bucket/profiler",
+        enable_ondemand_xprof=True,
+        run_name="test_run",
+        enable_profiler=False,
+        tensorboard_dir="/tmp/fake_tensorboard",
+    )
+
+    self.assertTrue(max_utils._ml_diagnostics_profiler_enabled(config))
+
+    with max_utils.Profiler(config, session_name="test_session"):
+      mock_xprof.return_value.start.assert_called_once_with("test_session")
+
+    mock_xprof.return_value.stop.assert_called_once()
+
+  @patch("jax.profiler.start_trace")
+  @patch("maxdiffusion.max_utils.machinelearning_run")
+  @patch("maxdiffusion.max_utils.xprof")
+  @patch("jax.process_index", return_value=0)
+  def test_both_profilers_enabled_prioritizes_mld(
+      self, mock_process_index, mock_xprof, mock_ml_run, mock_start_trace
+  ):
+    """Tests that when both ML Diagnostics and JAX profiler are enabled, ML Diagnostics is prioritized and JAX profiler is skipped."""
+    config = MockConfig(
+        enable_ml_diagnostics=True,
+        profiler_gcs_path="gs://fake-bucket/profiler",
+        enable_ondemand_xprof=True,
+        run_name="test_run",
+        enable_profiler=True,
+        tensorboard_dir="/tmp/fake_tensorboard",
+    )
+
+    with max_utils.Profiler(config, session_name="test_session"):
+      mock_xprof.return_value.start.assert_called_once_with("test_session")
+      mock_start_trace.assert_not_called()
+
+    mock_xprof.return_value.stop.assert_called_once()
+
   @patch("jax.profiler.start_trace")
   @patch("jax.profiler.stop_trace")
   @patch("jax.process_index", return_value=0)

src/maxdiffusion/train_utils.py

@@ -61,6 +61,20 @@ def _validate_gcs_bucket_name(bucket_name, config_var):
         " Set either train_text_encoder, or cache_latents_text_encoder_outputs to False"
     )
 
+  keys = config.get_keys()
+  if (
+      "enable_profiler" in keys
+      and config.enable_profiler
+      and "enable_ml_diagnostics" in keys
+      and config.enable_ml_diagnostics
+  ):
+    max_logging.log(
+        "Both enable_profiler and enable_ml_diagnostics are set to True. They"
+        " share the underlying JAX tracer and cannot run concurrently."
+        " Prioritizing ML Diagnostics profiler and disabling standard JAX"
+        " profiler."
+    )
+
 
 def record_scalar_metrics(metrics, step_time_delta, per_device_tflops, lr):
   """Records scalar metrics to be written to tensorboard"""