Add option to enable on-demand profiling server in MaxText ML Diagnostics

- Add `managed_mldiagnostics_on_demand_profiling` configuration flag to `base.yml` (enabled by default) and `types.py`.
- Pass `on_demand_xprof` parameter to `mldiag.machinelearning_run(...)` in `managed_mldiagnostics.py` to allow users to toggle on-demand profiling at runtime.

Fixes AI-Hypercomputer#4129.

PiperOrigin-RevId: 932303562

Author: Google-ML-Automation

src/maxtext/common/managed_mldiagnostics.py

@@ -72,5 +72,6 @@ def should_log_key(key, value):
         run_group=config.managed_mldiagnostics_run_group,
         configs=config_dict,
         gcs_path=config.managed_mldiagnostics_dir,
+        on_demand_xprof=config.managed_mldiagnostics_on_demand_profiling,
         region=region,
     )

src/maxtext/configs/base.yml

@@ -838,6 +838,7 @@ tpu_num_sparse_cores_to_trace: 2
 # - upload xplane profiling, if it is enabled.
 # - upload training metrics, at the defined log_period interval.
 managed_mldiagnostics: false # Whether to enable the managed diagnostics
+managed_mldiagnostics_on_demand_profiling: true # Enable on-demand profiling server by default
 managed_mldiagnostics_run_group: ""  # Optional. Used to group multiple runs.
 managed_mldiagnostics_region: "" # Optional. GCP region for managed mldiagnostics. If empty, it will be auto-detected by the SDK.
 

src/maxtext/configs/types.py

@@ -1785,6 +1785,10 @@ class ManagedMLDiagnostics(BaseModel):
   """Configuration for managed mldiagnostics."""
 
   managed_mldiagnostics: bool = Field(False, description="Enable managed mldiagnostics.")
+  managed_mldiagnostics_on_demand_profiling: bool = Field(
+      True,
+      description=("Enable on-demand profiling server when managed mldiagnostics is" " enabled."),
+  )
   managed_mldiagnostics_run_group: str = Field("", description="Name used to group multiple runs.")
   managed_mldiagnostics_region: str = Field("", description="GCP region for managed mldiagnostics.")