Support TPU advanced profiling options and upload_all_profiler_results in MicroStepProfiler

[SurbhiJainUSC]

Description

This PR aligns MicroStepProfiler in the training engine with common/profiler.py so it respects config.profiler, upload_all_profiler_results, and TPU/power advanced profiling options.

Previously, MicroStepProfiler:

• Did not check config.profiler, meaning ProfilerType.NONE (profiler="") still opened a trace whenever profiler_steps > 0 (which defaults to 5 in base.yml), causing "No profile started" failures under Pathways.
• Profiled only when jax.process_index() == 0, ignoring upload_all_profiler_results.
• Set ProfileOptions.advanced_configuration only for enable_continuous_profiling, ignoring enable_tpu_profiling_options (tpu_num_chips_to_profile_per_task, tpu_num_sparse_core_tiles_to_trace,
  tpu_num_sparse_cores_to_trace) and profile_power_events (tpu_power_trace_level, e2e_enable_fw_*_event).

This change updates MicroStepProfiler to:

• Gate profiling on config.profiler == ProfilerType.XPLANE.
• Allow non-zero hosts to profile when config.upload_all_profiler_results is True.
• Populate ProfileOptions.advanced_configuration for enable_tpu_profiling_options and profile_power_events in addition to enable_continuous_profiling.

FIXES: b/573295009

Tests

• Ran unit and integration tests for training engine profiling:

pytest tests/post_training/unit/maxtext_engine_profiling_test.py
pytest tests/post_training/integration/maxtext_engine_e2e_test.py

• Added unit tests in tests/post_training/unit/maxtext_engine_profiling_test.py

Checklist

Before submitting this PR, please make sure (put X in square brackets):

• I have performed a self-review of my code. For an optional AI review, add the gemini-review label.
• I have necessary comments in my code, particularly in hard-to-understand areas.
• I have run end-to-end tests tests and provided workload links above if applicable.
• I have made or will make corresponding changes to the doc if needed, including adding new documentation pages to the relevant Table of Contents (toctree directive) as explained in our documentation.

[gemini-code-assist]

Code Review

This pull request updates the MicroStepProfiler to conditionally enable profiling based on the profiler type, process index, and upload settings. It also dynamically populates the advanced_configuration dictionary with TPU profiling and power event options. Corresponding integration and unit tests have been added to verify these new behaviors. The review feedback suggests explicitly casting config.xprof_tpu_power_trace_level to an integer to avoid potential serialization issues with custom enum subclasses in the underlying profiling client.

[codecov]

Codecov Report

✅ All modified and coverable lines are covered by tests.

📢 Thoughts on this report? Let us know!