Wire MaxText RL profiling and metrics onto tunix's targeted RL profiler and ML Diagnostics.

[copybara-service]

Wire MaxText RL profiling and metrics onto tunix's targeted RL profiler and ML Diagnostics.

The Tunix-based RL trainer used tunix's SFT-era step profiler, which only
traces the actor on process 0 and never the vLLM rollout, and it had no
ML Diagnostics integration. Tunix now provides a per-role, invocation-counted
RLProfileConfig and an MLDiagScalarBackend; this change maps the MaxText
RL config onto them.

• profiler=xplane builds a tunix RLProfileConfig from the new
  rl.profiler_start_invocation, rl.profiler_num_invocations and
  rl.profiler_timeout_secs keys, producing one <role>_invocation_<n> trace
  per role. Non-managed traces land in tensorboard_dir; with
  managed_mldiagnostics=true the tunix ML Diagnostics backend writes under
  the MLRun's gcs_path (managed_mldiagnostics_dir), so
  RLProfileConfig.mldiagnostics_dir is not forwarded (tunix's JAX backend
  would otherwise prefer it over tensorboard_dir).
• managed_mldiagnostics=true creates the ML Diagnostics run before the
  RLCluster (tunix's RLProfiler preflights the active run), uploads the
  trainer config (plus sampler.<key> for any sampler-side differences on
  multi-slice topologies), and installs MLDiagScalarBackend next to
  TensorboardBackend as the tunix metrics backends. RLConfig now derives
  managed_mldiagnostics_dir like MaxTextConfig does and rejects a
  non-gs:// directory or on-demand profiling combined with profiler=xplane
  at config time.
• common/managed_mldiagnostics.py changes accordingly. ManagedMLDiagnostics
  stays a process-wide singleton (the RL trainer constructs it once at startup;
  pretrain's common/profiler.py is unchanged). It gains an optional
  sampler_config: when it is the same object as the trainer config
  (single-slice, colocated) the merge is skipped and the run config is a flat
  dict with no sampler. keys; when it is a distinct object (multi-slice,
  disaggregated) only keys whose values differ are added as sampler.<key>.
  The JSON-serializability filter now also catches ValueError so NaN/Inf
  values are dropped instead of failing the upload. Because
  gcloud_stub.mldiagnostics_modules() returns a stub object rather than
  None when the SDK is absent, availability is detected via the returned
  stub flag (mldiagnostics_available()); ManagedMLDiagnostics logs a
  warning when the SDK is in stub mode, while pretrain keeps its previous
  warn-and-continue behavior.
• With managed_mldiagnostics=false the metrics backends are left to tunix's
  defaults, so non-managed runs log exactly as before. With
  managed_mldiagnostics=true the backends are listed explicitly
  (MLDiagScalarBackend plus TensorboardBackend when enable_tensorboard
  is set); a non-empty explicit list replaces tunix's defaults, so managed runs
  do not get tunix's W&B backend in OSS (or the CLU backend internally).
• The step-window profiling flags (profiler_steps,
  skip_first_n_steps_for_profiler, profile_periodically_period, ...) are
  removed from RLConfig via a ProfilingCore mixin that only keeps
  profiler; they have no consumer on the Tunix RL path and are rejected with
  the list of valid fields. Pretrain/SFT keep the full Profiling mixin. The
  in-tree RL scripts are updated to the new keys.
• On Pathways, pathwaysutils traces worker 0 only (max_num_hosts=1), so a
  disaggregated sampler never appears in the rollout trace. A new
  rl.profiler_max_num_hosts key (default 1) parameterizes a wrapper around
  jax.profiler.start_trace to trace worker hosts across all slices.
• Format tfds_data_processing_c4_mlperf.py to satisfy repo-wide pyink pre-commit checks.

[codecov]

Codecov Report

❌ Patch coverage is 50.00000% with 33 lines in your changes missing coverage. Please review.

Files with missing lines	Patch %	Lines
src/maxtext/trainers/post_train/rl/train_rl.py	27.27%	28 Missing and 4 partials ⚠️
...t/input_pipeline/tfds_data_processing_c4_mlperf.py	0.00%	1 Missing ⚠️

📢 Thoughts on this report? Let us know!