Enhance Orbax save stats granularity with fine-grained duration metrics.

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,

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())
 

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)'

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: