From 668be1682458bca6eb1067e93f821251ba541969 Mon Sep 17 00:00:00 2001 From: Adam Cogdell Date: Tue, 28 Jul 2026 17:03:41 -0700 Subject: [PATCH] No public description PiperOrigin-RevId: 955551749 --- .../_src/checkpointers/async_checkpointer.py | 9 ++++++++ .../base_pytree_checkpoint_handler.py | 16 ++++++++++++++ .../orbax/checkpoint/_src/path/atomicity.py | 5 +++++ .../orbax/checkpoint/checkpoint_manager.py | 22 +++++++++++++++++++ 4 files changed, 52 insertions(+) diff --git a/checkpoint/orbax/checkpoint/_src/checkpointers/async_checkpointer.py b/checkpoint/orbax/checkpoint/_src/checkpointers/async_checkpointer.py index 5b8bc373a..314571b02 100644 --- a/checkpoint/orbax/checkpoint/_src/checkpointers/async_checkpointer.py +++ b/checkpoint/orbax/checkpoint/_src/checkpointers/async_checkpointer.py @@ -105,6 +105,10 @@ def _background_wait_for_commit_futures( '/jax/checkpoint/write/async/commit_duration_sec', commit_duration_secs, ) + jax.monitoring.record_event_duration_secs( + '/jax/orbax/write/async/tensorstore_duration_secs', + commit_duration_secs, + ) if process_count > 1: # All processes will wait at the barrier. When all processes are at the @@ -395,6 +399,7 @@ def _make_on_commit_callback( ) def _callback() -> None: + finalize_start_time = time.time() if utils.is_primary_host(self._primary_host): # Update StepMetadata after the handler save is complete. # (blocking write) @@ -433,6 +438,10 @@ def _callback() -> None: tmpdir, checkpoint_start_time, ) + jax.monitoring.record_event_duration_secs( + '/jax/orbax/write/async/finalize_duration_secs', + time.time() - finalize_start_time, + ) operation_recorder = event_tracking.OperationRecorder( tmpdir.get_final(), operation_type=event_tracking.OperationType.SAVE, diff --git a/checkpoint/orbax/checkpoint/_src/handlers/base_pytree_checkpoint_handler.py b/checkpoint/orbax/checkpoint/_src/handlers/base_pytree_checkpoint_handler.py index 8463b33cd..63cf26d0a 100644 --- a/checkpoint/orbax/checkpoint/_src/handlers/base_pytree_checkpoint_handler.py +++ b/checkpoint/orbax/checkpoint/_src/handlers/base_pytree_checkpoint_handler.py @@ -746,6 +746,14 @@ async def async_save( ), ) ] + jax.monitoring.record_event_duration_secs( + '/jax/orbax/write/async/tree_mapping_duration_secs', + batch_requests_ready_time - start_time, + ) + jax.monitoring.record_event_duration_secs( + '/jax/orbax/write/async/d2h_duration_secs', + total_serialization_initiated_time - batch_requests_ready_time, + ) async_save_end_time = time.time() logging.info( '[process=%s][thread=%s] Initiated Pytree async_save. Time taken:' @@ -1202,6 +1210,10 @@ async def _write_metadata_file( '/jax/checkpoint/write/async/metadata_write_duration_secs', time.time() - metadata_write_start_time, ) + jax.monitoring.record_event_duration_secs( + '/jax/orbax/write/async/metadata_write_duration_secs', + time.time() - metadata_write_start_time, + ) async def _write_metadata_after_commits( self, @@ -1370,6 +1382,10 @@ async def merge_ocdbt_per_process_files(): '/jax/checkpoint/write/async/ocdbt_merge_duration_secs', time.time() - merge_start_time, ) + jax.monitoring.record_event_duration_secs( + '/jax/orbax/write/async/ocdbt_merge_duration_secs', + time.time() - merge_start_time, + ) finalize_coros.append(merge_ocdbt_per_process_files()) diff --git a/checkpoint/orbax/checkpoint/_src/path/atomicity.py b/checkpoint/orbax/checkpoint/_src/path/atomicity.py index 35dc921c8..df88f2b4a 100644 --- a/checkpoint/orbax/checkpoint/_src/path/atomicity.py +++ b/checkpoint/orbax/checkpoint/_src/path/atomicity.py @@ -841,8 +841,13 @@ async def on_commit_callback( verification using set_immutable: Whether to mark all files as immutable. This is only """ + atomicity_start_time = time.time() await tmp_dir.finalize( ) + jax.monitoring.record_event_duration_secs( + '/jax/orbax/write/async/atomicity_duration_secs', + time.time() - atomicity_start_time, + ) record_saved_duration(checkpoint_start_time) logging.info( '[process=%s][thread=%s] Finished saving checkpoint (finalized tmp dir)' diff --git a/checkpoint/orbax/checkpoint/checkpoint_manager.py b/checkpoint/orbax/checkpoint/checkpoint_manager.py index 9a3083d09..991275ba3 100644 --- a/checkpoint/orbax/checkpoint/checkpoint_manager.py +++ b/checkpoint/orbax/checkpoint/checkpoint_manager.py @@ -1392,6 +1392,7 @@ def save( step_stats.step = step step_stats.checkpoint_manager_blocking_start_time = time.time() step_stats.directory = str(self.directory) + validation_start_time = time.time() if items is None and args is None: raise ValueError('Must provide `args` for `save`.') @@ -1507,6 +1508,17 @@ def save( logging.info( '[process=%s] Saving checkpoint at step %d', process_index, step ) + validation_duration = time.time() - validation_start_time + if is_async_checkpointer(self._checkpointer): + jax.monitoring.record_event_duration_secs( + '/jax/orbax/write/async/validation_duration_secs', + validation_duration, + ) + else: + jax.monitoring.record_event_duration_secs( + '/jax/orbax/write/validation_duration_secs', + validation_duration, + ) step_stats.checkpointer_blocking_start_time = time.time() self._checkpointer.save( save_directory, args=args, custom_metadata=custom_metadata, force=True @@ -2066,6 +2078,16 @@ def wait_until_finished(self): '/jax/checkpoint/write/wait_for_prev_duration_secs', duration, ) + if self._finalize_thread.get() is not None: + jax.monitoring.record_event_duration_secs( + '/jax/orbax/write/async/wait_for_prev_duration_secs', + duration, + ) + else: + jax.monitoring.record_event_duration_secs( + '/jax/orbax/write/wait_for_prev_duration_secs', + duration, + ) self._wait_for_prev_save_duration += duration def is_saving_in_progress(self) -> bool: