Description

Restore Orbax checkpoint Cloud Logging for ML Goodput after the Orbax V1 migration.

The V1 migration removed enable_checkpoint_cloud_logger: setup_checkpoint_logger and create_orbax_checkpoint_manager now only log a warning and drop the logger [cl/974668260]. This is because the Orbax V1 ocp.training.Checkpointer has no logger argument, and it builds its internal V0 CheckpointManager with the default StandardLogger, which doesn't write to Cloud Logging. As a result, checkpoint save/restore step statistics no longer reach the goodput_<run_name> log, and Goodput cannot measure checkpoint save/restore badput.

This PR:

• src/maxtext/common/checkpointing.py:
  • setup_checkpoint_logger: creates the Orbax CloudLogger again (log name goodput_<run_name>) when enable_checkpoint_cloud_logger=true. If CloudLogger isn't available, it warns and returns None.
  • create_orbax_checkpoint_manager: attaches the logger to the V1 Checkpointer's internal V0 manager, so V0-compatible save step statistics are logged again. This relies on a private attribute until Orbax V1 exposes a public logger option (TODO: b/529622681). If the internal attribute is missing, it warns instead of crashing.
  • load_state_if_possible: V1 load_checkpointables skips the internal manager's restore logging, so restore step statistics (RestoreStepStatistics) are logged around the load.
• src/maxtext/utils/train_utils.py: updates the outdated comment and removes a pylint suppression that no longer applies.

BUGS: b/568044767

Tests

Tested on TPU v6e-8 with enable_goodput_recording=true monitor_goodput=true enable_checkpoint_cloud_logger=true steps=20 checkpoint_period=5, then resumed the same run with steps=25:

• Save entries (event_type="save", steps 0/5/10/15/19) with checkpoint_manager_blocking_start_time / checkpoint_manager_blocking_duration_secs are back in the Goodput log: Log Explorer (save)
• In comparison, this run does not have this before this PR.
• On resume, the restore from step 19 (event_type="restore") is logged, followed by saves at steps 20/24: Log Explorer (restore).

Note: with checkpoint entries back, ml-goodput-measurement 0.2.3 reports Productive training time is invalid on this short async-checkpoint run (same symptom as b/553561110). That is a calculator-side issue, tracked separately.

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 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.