feat(goodput): use get_goodput_job_name in setup_checkpoint_logger

src/maxtext/common/checkpointing.py

@@ -36,6 +36,7 @@
 from maxtext.checkpoint_conversion.utils import load_dynamic
 from maxtext.common import checkpoint_context
 from maxtext.common import emergency_checkpointing
+from maxtext.common.goodput import get_goodput_job_name
 from maxtext.common import grain_utility
 from maxtext.common import train_state_nnx
 from maxtext.input_pipeline import multihost_dataloading
@@ -885,7 +886,8 @@ def setup_checkpoint_logger(config) -> Any | None:  # pytype: disable=attribute-
     config: MaxText config.
 
   Returns:
-    An Orbax `CloudLogger` writing to the `goodput_<run_name>` log if
+    An Orbax `CloudLogger` writing to the `goodput_<job_name>` log (where
+    `job_name` is resolved via `get_goodput_job_name(config)`) if
     `enable_checkpoint_cloud_logger` is set, otherwise None.
   """
   if not config.enable_checkpoint_cloud_logger:
@@ -899,11 +901,10 @@ def setup_checkpoint_logger(config) -> Any | None:  # pytype: disable=attribute-
         " (is google-cloud-logging installed?). Checkpoint logs are disabled."
     )
     return None
-  logger_name = f"goodput_{config.run_name}"
+  job_name = get_goodput_job_name(config)
+  logger_name = f"goodput_{job_name}"
   try:
-    orbax_cloud_logger = cloud_logger_cls(
-        options=cloud_logger_options_cls(job_name=config.run_name, logger_name=logger_name)
-    )
+    orbax_cloud_logger = cloud_logger_cls(options=cloud_logger_options_cls(job_name=job_name, logger_name=logger_name))
   except Exception as e:  # pylint: disable=broad-exception-caught
     max_logging.warning(f"Failed to initialize Orbax CloudLogger: {e}. Checkpoint logs are disabled.")
     return None

src/maxtext/common/goodput.py

@@ -29,7 +29,7 @@
 
 
 def get_goodput_job_name(config) -> str:
-  return config.goodput_job_name or config.run_name
+  return getattr(config, "goodput_job_name", "") or config.run_name
 
 
 class GoodputEvent(Enum):

src/maxtext/configs/base.yml

@@ -1230,7 +1230,7 @@ abort_on_inf_loss: true # Check for Inf and abort if found in training loss
 
 # Goodput parameters
 enable_goodput_recording: false
-goodput_job_name: ""
+goodput_job_name: "" # Optional/test only; not needed for production runs (defaults to run_name)
 monitor_goodput: false
 goodput_upload_interval_seconds: 30
 enable_pathways_goodput: false

src/maxtext/configs/post_train/rl.yml

@@ -106,7 +106,7 @@ debug: false
 # uploaded to tensorboard.
 enable_tunix_perf_metrics: false
 enable_goodput_recording: false
-goodput_job_name: ""
+goodput_job_name: "" # Optional/test only; not needed for production runs (defaults to run_name)
 monitor_goodput: false
 enable_pathways_goodput: false
 

tests/unit/goodput_utils_test.py

@@ -572,6 +572,19 @@ def test_setup_checkpoint_logger_flag_on_uses_goodput_log(self):
     self.assertIs(logger, cloud_logger.return_value)
     options.assert_called_once_with(job_name="run", logger_name="goodput_run")
 
+  def test_setup_checkpoint_logger_uses_goodput_job_name_override(self):
+    config = SimpleNamespace(
+        enable_checkpoint_cloud_logger=True, run_name="shared_run", goodput_job_name="gemma3-4b-pre-shared_run"
+    )
+    with (
+        mock.patch.object(checkpointing.ocp_logging, "CloudLogger", create=True) as cloud_logger,
+        mock.patch.object(checkpointing.ocp_logging, "CloudLoggerOptions", create=True) as options,
+    ):
+      logger = checkpointing.setup_checkpoint_logger(config)
+
+    self.assertIs(logger, cloud_logger.return_value)
+    options.assert_called_once_with(job_name="gemma3-4b-pre-shared_run", logger_name="goodput_gemma3-4b-pre-shared_run")
+
 
 if __name__ == "__main__":
   unittest.main()