feat(goodput): add goodput_job_name config and instrument RL & Distillation trainers

[KevinYenky]

Description

This PR extends MaxText's Goodput measurement and monitoring integration on post-training workflows so that end-to-end Airflow/Cloud Composer DAGs can uniquely track and validate Goodput per workload:

1. Configurable goodput_job_name:

• Adds goodput_job_name (default "") to post_train/rl.yml.
• Adds enable_goodput_recording, goodput_job_name, monitor_goodput, and enable_pathways_goodput defaults to src/maxtext/configs/post_train/rl.yml (which does not inherit base.yml) so M_* environment variable overrides in pyconfig are applied to RL workloads.

2. RL Trainer & Hooks Goodput Instrumentation (train_rl.py, rl/hooks.py, rl/utils_rl.py):

• Instruments rl_train in src/maxtext/trainers/post_train/rl/train_rl.py with create_goodput_recorder, maybe_monitor_goodput, RECORD_JOB_START_TIME, GoodputEvent.TPU_INIT, GoodputEvent.TRAINING_PREPARATION, and RECORD_JOB_END_TIME (recorded in a finally block).
• Updates RLTrainingHooks (src/maxtext/trainers/post_train/rl/hooks.py) and install_training_hooks (src/maxtext/trainers/post_train/rl/utils_rl.py) to accept goodput_recorder, record record_step_start_time(train_ctx.train_steps) on on_train_step_start, and safely guard on_train_step_end when eval_interval <= 0 or num_test_batches <= 0 so hooks can be installed for Goodput-only runs without intermediate evaluation.

3. Distillation Trainer & Shared BaseTrainingHooks Instrumentation (train_distill.py, post_train/hooks.py):

• Instruments train_distill in src/maxtext/trainers/post_train/distillation/train_distill.py with create_goodput_recorder,
  maybe_monitor_goodput, RECORD_JOB_START_TIME, GoodputEvent.TPU_INIT, GoodputEvent.TRAINING_PREPARATION, BaseTrainingHooks(student_config, goodput_recorder), and RECORD_JOB_END_TIME (recorded in a finally block).
• Updates BaseTrainingHooks in src/maxtext/trainers/post_train/hooks.py so mesh and learning_rate_schedule are optional, enabling Goodput-only hook instantiation (BaseTrainingHooks(config, goodput_recorder)) while remaining backward-compatible with SFTTrainingHooks and DPOTrainingHooks.

4. Post-Training Dependencies (tpu-post-train-requirements.txt):

• Adds ml-goodput-measurement>=0.2.2 to src/dependencies/requirements/base_requirements/tpu-post-train-requirements.txt so post-training Docker images include the Goodput measurement library.

Tests

1. Unit Tests:

• Added unit tests for goodput_job_name fallback and override behavior in tests/unit/goodput_utils_test.py.
• Added RLGoodputTest in tests/post_training/unit/train_rl_test.py covering rl_train lifecycle/step recording (enable_goodput_recording=True and False) and RLTrainingHooks with eval_interval=0.
• Added test_train_distill_goodput_enabled and test_train_distill_goodput_disabled in tests/post_training/unit/train_distill_test.py.
• Added test_base_training_hooks_goodput_only_mode in tests/post_training/unit/hooks_test.py.

pytest \
tests/unit/goodput_utils_test.py \
tests/post_training/unit/train_rl_test.py \
tests/post_training/unit/train_distill_test.py \
tests/post_training/unit/hooks_test.py \
tests/post_training/unit/sft_hooks_test.py \
tests/post_training/unit/dpo_hooks_test.py

2. End-to-End Composer DAG Verification:
   • Verified and post-training (SFT, LoRA, Multimodal SFT, RL) Goodput recording in Cloud Logging via maxtext_e2e_tpu_post_training.

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 integrates goodput measurement and recording into both distillation and reinforcement learning (RL) post-training pipelines, adding configuration options, lifecycle hooks, and corresponding unit tests. The review feedback suggests improving robustness and code cleanliness by using the defensive 'record_goodput' helper to avoid potential AttributeError exceptions, and passing 'goodput_recorder' as an explicit keyword argument to 'BaseTrainingHooks' rather than relying on fragile positional argument shifting and type-checking in the constructor.

[codecov]

Codecov Report

❌ Patch coverage is 90.08264% with 12 lines in your changes missing coverage. Please review.

Files with missing lines	Patch %	Lines
.../trainers/post_train/distillation/train_distill.py	85.29%	8 Missing and 2 partials ⚠️
src/maxtext/trainers/post_train/rl/hooks.py	83.33%	0 Missing and 1 partial ⚠️
src/maxtext/trainers/post_train/rl/train_rl.py	93.75%	0 Missing and 1 partial ⚠️

📢 Thoughts on this report? Let us know!

[dipannita08]

This is not blocking but I just realized that checkpoint logs are turned off in the pre-train DAGs.
We can enable them in follow up PR, by doing two things:

1. setup_checkpoint_logger() in checkpointing.py should construct job_name = get_goodput_job_name(config) and logger_name = f"goodput_{job_name}", so that this works for prod as well as these DAGs.
2. M_ENABLE_CHECKPOINT_CLOUD_LOGGER=true in maxtext_e2e_tpu_pre_training.py

Don't block this PR, it is completely safe to have Checkpoint badput = 0 until we enable it in a follow up.

[dipannita08]

BTW, small nit but feel free to add a comment here saying "optional/test only" so users know they do not need to set it for production runs.

[dipannita08]

Left a small change required as a follow up PR, otherwise LGTM. Thank you!