Ensure AsyncCheckpointer completion logs on every host instead of jus…

…t the leader.

PiperOrigin-RevId: 724466022

checkpoint/CHANGELOG.md

@@ -7,10 +7,17 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 
 ## [Unreleased]
 
+## [0.11.4] - 2025-02-07
+
 ### Changed
 
 - Updated orbax-checkpoint PyPI package to exclude tests.
 
+### Fixed
+
+- `AsyncCheckpointer` completion logging, to log on all hosts instead of just
+the leader.
+
 ## [0.11.3] - 2025-02-06
 
 ### Changed

checkpoint/orbax/checkpoint/_src/checkpointers/async_checkpointer.py

@@ -58,11 +58,6 @@ def _on_commit_callback(
       '/jax/checkpoint/write/async/total_duration_secs',
       total_duration_secs,
   )
-  logging.info(
-      'Finished asynchronous save in %.2f seconds to %s',
-      total_duration_secs,
-      tmpdir.get_final(),
-  )
 
 
 def _add_deadline_exceeded_notes(e: jax.errors.JaxRuntimeError):
@@ -405,7 +400,8 @@ def _callback() -> None:
         # Update StepMetadata after the handler save is complete.
         # (blocking write)
         self._save_step_metadata(tmpdir.get(), custom_metadata=custom_metadata)
-      logging.info(
+      logging.vlog(
+          1,
           '[process=%s][thread=%s] Async Save Callback [1/3]: Finalizing'
           ' Handler: %s on %s',
           multihost.process_index(),
@@ -415,7 +411,8 @@ def _callback() -> None:
       )
       # Finalize does a final StepMetadata update.
       self._handler.finalize(tmpdir.get())
-      logging.info(
+      logging.vlog(
+          1,
           '[process=%s][thread=%s] Async Save Callback [2/3]: Running'
           ' post_finalization_callback: %s on %s',
           multihost.process_index(),
@@ -425,7 +422,8 @@ def _callback() -> None:
       )
       if self._post_finalization_callback is not None:
         self._post_finalization_callback()
-      logging.info(
+      logging.vlog(
+          1,
           '[process=%s][thread=%s] Async Save Callback [3/3]: Finalizing'
           ' checkpoint directory: %s',
           multihost.process_index(),
@@ -436,6 +434,12 @@ def _callback() -> None:
           tmpdir,
           checkpoint_start_time,
       )
+      logging.info(
+          'Finished asynchronous save (blocking + background) in %.2f seconds'
+          ' to %s',
+          time.time() - checkpoint_start_time,
+          directory,
+      )
 
     self._async_manager.start_async_commit(
         directory,

checkpoint/orbax/checkpoint/version.py

@@ -16,7 +16,7 @@
 
 # A new PyPI release will be pushed everytime `__version__` is increased.
 # Also modify version and date in CHANGELOG.
-__version__ = '0.11.3'
+__version__ = '0.11.4'
 
 
 # TODO: b/362813406 - Add latest change timestamp and commit number.