Add Orbax checkpoint logger for AXLearn

Author: dipannita08

axlearn/cloud/gcp/measurement.py

@@ -25,6 +25,7 @@
 from typing import Optional, Sequence
 
 import jax
+import orbax.checkpoint as ocp
 from absl import flags, logging
 from ml_goodput_measurement import goodput
 from ml_goodput_measurement import monitoring as goodput_monitoring
@@ -134,6 +135,19 @@ def record_event(self, event: measurement.EventType, *args, **kwargs):
                     )
         # pylint: enable=try-except-raise
 
+    def create_checkpoint_logger(self) -> Optional[ocp.logging.CloudLogger]:
+        try:
+            logging.info("Creating a Goodput checkpoint logger.")
+            return ocp.logging.CloudLogger(
+                options=ocp.logging.CloudLoggerOptions(
+                    job_name=self._job_name,
+                    logger_name=self._logger_name,
+                )
+            )
+        except Exception as e:  # pylint: disable=broad-exception-caught
+            logging.warning("Failed to create Goodput checkpoint logger: %s", e, exc_info=True)
+            return None
+
     @contextlib.contextmanager
     def _maybe_monitor_goodput(self, *args, **kwargs):
         """Monitor cumulative goodput if enabled.

axlearn/cloud/gcp/measurement_test.py

@@ -373,3 +373,46 @@ def test_maybe_monitor_all(
             else:
                 mock_monitor_instance.start_rolling_window_goodput_uploader.assert_not_called()
                 mock_monitor_instance.stop_rolling_window_goodput_uploader.assert_not_called()
+
+    @mock.patch("jax.process_index", return_value=0)
+    def test_create_checkpoint_logger_success(self, _):
+        """Tests that create_checkpoint_logger creates a CloudLogger with correct config."""
+        cfg = GoodputRecorder.default_config().set(
+            name="test-job",
+            upload_dir="/test",
+            upload_interval=30,
+        )
+        recorder = GoodputRecorder(cfg)
+
+        with mock.patch("orbax.checkpoint.logging.CloudLogger") as mock_logger_cls:
+            mock_logger_instance = mock_logger_cls.return_value
+            logger = recorder.create_checkpoint_logger()
+
+            mock_logger_cls.assert_called_once()
+            self.assertIs(logger, mock_logger_instance)
+
+            _, kwargs = mock_logger_cls.call_args
+            options = kwargs["options"]
+            self.assertEqual(options.job_name, "test-job")
+            self.assertEqual(options.logger_name, "goodput_logger_test-job")
+
+    @mock.patch("jax.process_index", return_value=0)
+    def test_create_checkpoint_logger_failure(self, _):
+        """Tests that create_checkpoint_logger logs a warning on failure and returns None."""
+        cfg = GoodputRecorder.default_config().set(
+            name="fail-job",
+            upload_dir="/test",
+            upload_interval=30,
+        )
+        recorder = GoodputRecorder(cfg)
+
+        with mock.patch(
+            "orbax.checkpoint.logging.CloudLogger", side_effect=RuntimeError("TestError")
+        ) as mock_logger_cls, mock.patch.object(logging, "warning") as mock_warning:
+            logger = recorder.create_checkpoint_logger()
+            self.assertIsNone(logger)
+            mock_logger_cls.assert_called_once()
+            mock_warning.assert_called_once()
+            self.assertIn(
+                "Failed to create Goodput checkpoint logger", mock_warning.call_args[0][0]
+            )

axlearn/common/checkpointer_orbax.py

@@ -17,7 +17,7 @@
 import tensorflow as tf
 from absl import logging
 
-from axlearn.common import utils
+from axlearn.common import measurement, utils
 from axlearn.common.checkpointer import (
     STEP_NUM_DIGITS,
     STEP_PREFIX,
@@ -232,6 +232,9 @@ def save_fn_with_summaries(step: int, last_saved_step: Optional[int]) -> bool:
             step_prefix=STEP_PREFIX,
             step_format_fixed_length=STEP_NUM_DIGITS,
         )
+        self._checkpoint_logger = None
+        if measurement.global_recorder:
+            self._checkpoint_logger = measurement.global_recorder.create_checkpoint_logger()
         self._manager = ocp.CheckpointManager(
             directory=cfg.dir,
             options=ocp.CheckpointManagerOptions(
@@ -255,6 +258,7 @@ def save_fn_with_summaries(step: int, last_saved_step: Optional[int]) -> bool:
                     restore_concurrent_gb=cfg.max_concurrent_restore_gb,
                 ),
             },
+            logger=self._checkpoint_logger,
         )
 
     def _get_spec(self, *, step: int, state: Nested[Any]) -> Nested[Any]:

axlearn/common/checkpointer_orbax_emergency.py

@@ -27,7 +27,7 @@
 from jax.experimental.array_serialization import serialization
 
 from axlearn.common import file_system as fs
-from axlearn.common import utils, utils_spmd
+from axlearn.common import measurement, utils, utils_spmd
 from axlearn.common.checkpointer import (
     STEP_NUM_DIGITS,
     STEP_PREFIX,
@@ -667,6 +667,9 @@ def _composite_save_policy(*, step: int, evaler_summaries: dict[str, Any]):
         # See comments of _eval_summaries in `OrbaxCheckpointer`.
         self._eval_summaries = None
         self._reached_preemption = False
+        self._checkpoint_logger = None
+        if measurement.global_recorder:
+            self._checkpoint_logger = measurement.global_recorder.create_checkpoint_logger()
 
     # pylint: disable-next=redefined-builtin
     def ckpt_dir(self, step: int, dir: Optional[str] = None) -> str:
@@ -731,6 +734,7 @@ def _orbax_save_fn(
                 cleanup_tmp_directories=True,
                 enable_async_checkpointing=True,
             ),
+            logger=self._checkpoint_logger,
         )
         return self._tensor_manager
 

axlearn/common/checkpointer_orbax_emergency_test.py

@@ -10,6 +10,7 @@
 import tempfile
 from contextlib import ExitStack, closing
 from typing import Optional
+from unittest import mock
 
 import jax
 import numpy as np
@@ -18,7 +19,7 @@
 from absl.testing import parameterized
 from jax import numpy as jnp
 
-from axlearn.common import utils_spmd
+from axlearn.common import measurement, utils_spmd
 from axlearn.common.checkpointer_orbax_emergency import (
     OrbaxEmergencyCheckpointer,
     _dump_process_info,
@@ -299,3 +300,28 @@ def start_processes(reverse_process_id: bool = False):
             finally:
                 for p in processes:
                     p.kill()
+
+    @mock.patch("orbax.checkpoint._src.multihost.multihost.initialize_runtime_to_distributed_ids")
+    @mock.patch("orbax.checkpoint._src.multihost.multihost.initialize_distributed_to_device_ids")
+    def test_emergency_checkpointer_initializes_logger_from_global_recorder(
+        self, mock_init_runtime, mock_init_device_ids
+    ):  # pylint: disable=unused-argument
+        """Tests OrbaxEmergencyCheckpointer initializes _checkpoint_logger."""
+        with tempfile.TemporaryDirectory() as temp_dir, mock.patch.object(
+            measurement, "global_recorder", mock.MagicMock()
+        ) as mock_recorder:
+            mock_logger = mock.MagicMock()
+            mock_recorder.create_checkpoint_logger.return_value = mock_logger
+
+            cfg = OrbaxEmergencyCheckpointer.default_config().set(
+                name="test_logger",
+                trainer_dir=temp_dir,
+                dir=temp_dir,
+                local_dir=temp_dir,
+                replica_axis_index=0,
+            )
+
+            ckpt: OrbaxEmergencyCheckpointer = cfg.instantiate(parent=None)
+
+            mock_recorder.create_checkpoint_logger.assert_called_once()
+            self.assertEqual(ckpt._checkpoint_logger, mock_logger)

axlearn/common/checkpointer_orbax_test.py

@@ -10,13 +10,14 @@
 import os
 import tempfile
 from typing import Sequence
+from unittest import mock
 
 import jax
 import orbax.checkpoint as ocp
 from jax import numpy as jnp
 from jax.experimental import mesh_utils
 
-from axlearn.common import test_utils
+from axlearn.common import measurement, test_utils
 from axlearn.common.checkpointer import read_index_file
 from axlearn.common.checkpointer_orbax import OrbaxCheckpointer
 
@@ -52,3 +53,21 @@ def test_index(self):
                 ),
             )
             self.assertEqual(ref_index, test_index["index"])
+
+    def test_initializes_checkpoint_logger_from_global_recorder(self):
+        """Tests that OrbaxCheckpointer initializes _checkpoint_logger if global_recorder is set."""
+        with tempfile.TemporaryDirectory() as temp_dir, mock.patch.object(
+            measurement, "global_recorder", mock.MagicMock()
+        ) as mock_recorder:
+            mock_logger = mock.MagicMock(spec=ocp.logging.CloudLogger)
+            mock_recorder.create_checkpoint_logger.return_value = mock_logger
+
+            ckpt = (
+                OrbaxCheckpointer.default_config()
+                .set(name="test", dir=temp_dir)
+                .instantiate(parent=None)
+            )
+
+            # Ensure create_checkpoint_logger was called and the logger was set.
+            mock_recorder.create_checkpoint_logger.assert_called_once()
+            self.assertEqual(ckpt._checkpoint_logger, mock_logger)

axlearn/common/measurement.py

@@ -102,6 +102,10 @@ def record_event(self, event: Event, *args, **kwargs):
     def maybe_monitor_all(self):
         yield
 
+    def create_checkpoint_logger(self) -> Optional[object]:
+        """Optionally returns a fully functional and independent checkpoint logger."""
+        return None
+
 
 _recorders: dict[str, type] = {}
 _T = TypeVar("_T")

axlearn/common/measurement_test.py

@@ -96,3 +96,6 @@ def test_initialize(self, recorder_type, expected):
         # Ensure that maybe_monitor_all does not fail (just enter and exit context).
         with measurement.global_recorder.maybe_monitor_all():
             pass
+
+        # Ensure that create_checkpoint_logger does not crash.
+        measurement.global_recorder.create_checkpoint_logger()