Enable power collection in xprof MaxText.

This change adds configuration options to enable TPU power tracing and firmware event collection (throttle, power level, thermal) in the XPlane profiler.

PiperOrigin-RevId: 867809819

Author: Google-ML-Automation

src/maxtext/common/profiler.py

@@ -50,6 +50,15 @@ def __init__(self, config, offset_step=0):
     if config.managed_mldiagnostics:
       ManagedMLDiagnostics(config)  # Initialize the MLRun instance.
 
+    self.profiling_options = jax.profiler.ProfileOptions()
+    if self.mode == "xplane" and not self.managed_mldiagnostics:
+      self.profiling_options.advanced_configuration = {
+          "tpu_power_trace_level": config.xprof_tpu_power_trace_level,
+          "e2e_enable_fw_throttle_event": config.xprof_e2e_enable_fw_throttle_event,
+          "e2e_enable_fw_power_level_event": config.xprof_e2e_enable_fw_power_level_event,
+          "e2e_enable_fw_thermal_event": config.xprof_e2e_enable_fw_thermal_event,
+      }
+
   def maybe_activate_profiler(self, step, state):
     """Conditionally activates the profiler based on the current step.
     This method checks if the current training step matches the step designated
@@ -87,7 +96,7 @@ def activate(self, blocking_object=None, optional_postfix=""):
         return
       self.libcudart.cudaProfilerStart()
     elif self.mode == "xplane":
-      jax.profiler.start_trace(self.output_path)
+      jax.profiler.start_trace(self.output_path, profiler_options=self.profiling_options)
 
   def maybe_deactivate_profiler(self, step, state):
     """Conditionally deactivates the profiler based on the current step.

src/maxtext/configs/base.yml

@@ -880,6 +880,12 @@ prometheus_port: 0
 enable_jax_profiler: False
 jax_profiler_port: 9999
 
+# TPU power trace level for xprof. 0:POWER_TRACE_NONE, 1:POWER_TRACE_NORMAL, or 2:POWER_TRACE_SPI
+xprof_tpu_power_trace_level: 0
+xprof_e2e_enable_fw_throttle_event: False
+xprof_e2e_enable_fw_power_level_event: False
+xprof_e2e_enable_fw_thermal_event: False
+
 log_config: True # Prints the config (after defaults have been set by pyconfig logic)
 debug_sharding: False # Prints model weights sharding info
 

src/maxtext/configs/types.py

@@ -16,28 +16,36 @@
 
 # pylint: disable=too-many-lines
 
-from enum import Enum
-from math import prod
-from tempfile import gettempdir
-from typing import Any, NewType, Literal, Optional
 import datetime
+import enum
+from enum import Enum
 import logging
 import math
+from math import prod
 import os
+from tempfile import gettempdir
+from typing import Any, Literal, NewType, Optional
 
 import jax
-
-from pydantic.config import ConfigDict
-from pydantic.fields import Field
-from pydantic.functional_validators import model_validator, field_validator
-from pydantic.main import BaseModel
-from pydantic.types import PositiveInt, NonNegativeFloat, NonNegativeInt
-
 from MaxText import accelerator_to_spec_map
 from MaxText.common_types import AttentionType, DecoderBlockType, ShardMode
 from MaxText.globals import MAXTEXT_ASSETS_ROOT
 from maxtext.utils import gcs_utils
 from maxtext.utils import max_utils
+from pydantic.config import ConfigDict
+from pydantic.fields import Field
+from pydantic.functional_validators import field_validator, model_validator
+from pydantic.main import BaseModel
+from pydantic.types import NonNegativeFloat, NonNegativeInt, PositiveInt
+
+
+class XProfTPUPowerTraceMode(enum.IntEnum):  # pylint: disable=invalid-name
+  """Enum for XProfTPUPowerTraceMode."""
+
+  POWER_TRACE_NONE = 0
+  POWER_TRACE_NORMAL = 1
+  POWER_TRACE_SPI = 2
+
 
 logger = logging.getLogger(__name__)
 
@@ -1316,6 +1324,16 @@ class Profiling(BaseModel):
   hide_profiler_step_metric: bool = Field(False, description="Whether to enable profiler step metric.")
   enable_jax_profiler: bool = Field(False, description="Enable the JAX live profiler.")
   jax_profiler_port: int = Field(9999, description="Port for the JAX profiler.")
+  xprof_tpu_power_trace_level: XProfTPUPowerTraceMode = Field(
+      XProfTPUPowerTraceMode.POWER_TRACE_NONE,
+      description=(
+          "TPU power trace level. The value should be 0 (POWER_TRACE_NONE), 1"
+          " (POWER_TRACE_NORMAL), or 2 (POWER_TRACE_SPI)"
+      ),
+  )
+  xprof_e2e_enable_fw_throttle_event: bool = Field(False, description="Enable FW throttle event.")
+  xprof_e2e_enable_fw_power_level_event: bool = Field(False, description="Enable FW power level event.")
+  xprof_e2e_enable_fw_thermal_event: bool = Field(False, description="Enable FW thermal event.")
 
 
 class HloDump(BaseModel):

tests/unit/profiler_test.py

@@ -13,19 +13,99 @@
 # limitations under the License.
 
 """Profiler tests."""
+import os
 import sys
 import unittest
+from unittest.mock import MagicMock, patch
 
 import pytest
 
 from MaxText import pyconfig
+from MaxText.globals import MAXTEXT_PKG_DIR
+from maxtext.configs import types
 from maxtext.common import profiler
 from tests.utils.test_helpers import get_test_config_path
 
 
 class ProfilerTest(unittest.TestCase):
   """Test for profiler."""
 
+  def setUp(self):
+    super().setUp()
+    # Mock jax.devices() to be deterministic and avoid runtime initialization errors
+    self.mock_devices = [MagicMock(slice_index=0) for _ in range(1)]
+    self.jax_patcher = patch("jax.devices", return_value=self.mock_devices)
+    self.jax_patcher.start()
+    self.jax_process_index_patcher = patch("jax.process_index", return_value=0)
+    self.jax_process_index_patcher.start()
+
+  def tearDown(self):
+    self.jax_patcher.stop()
+    self.jax_process_index_patcher.stop()
+    super().tearDown()
+
+  @pytest.mark.tpu_only
+  def test_profiler_options_populated_from_config(self):
+    """Verifies that Profiler initializes jax.profiler.ProfileOptions from config."""
+    config = pyconfig.initialize(
+        [sys.argv[0], os.path.join(MAXTEXT_PKG_DIR, "configs", "base.yml")],
+        enable_checkpointing=False,
+        run_name="test_profiler_options",
+        base_output_directory="/tmp",
+        profiler="xplane",
+        xprof_tpu_power_trace_level=1,
+        xprof_e2e_enable_fw_throttle_event=True,
+        xprof_e2e_enable_fw_power_level_event=True,
+        xprof_e2e_enable_fw_thermal_event=True,
+    )
+
+    with patch("jax.profiler.ProfileOptions") as mock_options_cls:
+      # Setup mock return value
+      mock_options_instance = MagicMock()
+      mock_options_cls.return_value = mock_options_instance
+
+      prof = profiler.Profiler(config)
+
+      # Check if ProfileOptions was instantiated
+      mock_options_cls.assert_called_once()
+
+      # Verify advanced_configuration was populated
+      expected_advanced_config = {
+          "tpu_power_trace_level": types.XProfTPUPowerTraceMode.POWER_TRACE_NORMAL,
+          "e2e_enable_fw_throttle_event": True,
+          "e2e_enable_fw_power_level_event": True,
+          "e2e_enable_fw_thermal_event": True,
+      }
+      self.assertEqual(prof.profiling_options.advanced_configuration, expected_advanced_config)
+
+  @pytest.mark.tpu_only
+  @patch("jax.profiler.start_trace")
+  def test_profiler_activate_passes_options(self, mock_start_trace):
+    """Verifies that activate() passes the profiling_options to jax.profiler.start_trace."""
+    config = pyconfig.initialize(
+        [sys.argv[0], os.path.join(MAXTEXT_PKG_DIR, "configs", "base.yml")],
+        enable_checkpointing=False,
+        run_name="test_profiler_options",
+        base_output_directory="/tmp",
+        profiler="xplane",
+        xprof_tpu_power_trace_level=2,
+    )
+
+    # We need to mock ProfileOptions as well to check identity or value
+    with patch("jax.profiler.ProfileOptions"):
+      prof = profiler.Profiler(config)
+      prof.activate()
+
+      # Verify start_trace was called with profiler_options
+      mock_start_trace.assert_called_once()
+      _, kwargs = mock_start_trace.call_args
+      self.assertIn("profiler_options", kwargs)
+      self.assertEqual(kwargs["profiler_options"], prof.profiling_options)
+      self.assertEqual(
+          prof.profiling_options.advanced_configuration["tpu_power_trace_level"],
+          types.XProfTPUPowerTraceMode.POWER_TRACE_SPI,
+      )
+
   # These periodic proilfer tests can run on any platform (cpu, gpu or tpu)
   @pytest.mark.tpu_only
   def test_periodic_profiler_third_period_starts(self):