feat(goodput): use get_goodput_job_name in setup_checkpoint_logger

[KevinYenky]

Description

Follow-up to PR#5563 to ensure Orbax checkpoint cloud logs are routed to the same Goodput logger bucket (goodput_<goodput_job_name>) when goodput_job_name is set.

Key Changes

1. src/maxtext/common/checkpointing.py:

• Updated setup_checkpoint_logger(config) to resolve job_name = get_goodput_job_name(config) and logger_name = f"goodput_{job_name}" instead of hardcoding config.run_name, so checkpoint Badput (UNPRODUCTIVE_CHECKPOINT_SAVE_TIME / UNPRODUCTIVE_CHECKPOINT_RESTORE_TIME) is logged to the per-workload bucket in multi-model E2E DAGs as well as production runs.

2. src/maxtext/common/goodput.py:

• Updated get_goodput_job_name(config) to safely fall back when goodput_job_name is absent on lightweight test mocks(getattr(config, "goodput_job_name", "") or config.run_name).

3. src/maxtext/configs/base.yml & src/maxtext/configs/post_train/rl.yml:

• Added inline comments clarifying that goodput_job_name is optional/test-only and defaults to run_name for production runs.

4. tests/unit/goodput_utils_test.py:

• Added unit test test_setup_checkpoint_logger_uses_goodput_job_name_override.

Tests

pytest tests/unit/goodput_utils_test.py -k "setup_checkpoint_logger or goodput_job_name" -v

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 checkpoint logger setup to use the get_goodput_job_name helper function, which safely retrieves the goodput_job_name attribute from the configuration (falling back to run_name if not present). It also adds goodput configuration parameters to the RL post-training configuration (rl.yml), documents the optional nature of goodput_job_name in base.yml, and adds a unit test to verify the job name override behavior. There are no review comments, and I have no additional feedback to provide.

[codecov]

Codecov Report

✅ All modified and coverable lines are covered by tests.

📢 Thoughts on this report? Let us know!