Refactor benchmark metrics to support multiple measurements per operation without needing context manager for each one.

Also,
*   Adding `IOBytesMetric` to track I/O read/write bytes and throughput.
*   Introducing a `METRIC_REGISTRY` to map metric keys to classes.

PiperOrigin-RevId: 831668357

Author: mridul-sahu

checkpoint/orbax/checkpoint/_src/testing/benchmarks/array_handler_benchmark.py

@@ -174,7 +174,7 @@ def test_fn(self, test_context: core.TestContext) -> core.TestResult:
     save_args = type_handlers.SaveArgs()
 
     # --- Serialization ---
-    with metrics.time('serialize'):
+    with metrics.measure('serialize'):
 
       async def serialize_and_wait():
         serialize_futures = await handler.serialize(
@@ -191,7 +191,7 @@ async def serialize_and_wait():
     logging.info('Serialization complete for %s', param_info.name)
 
     if options.use_ocdbt:
-      with metrics.time('merge_ocdbt'):
+      with metrics.measure('merge_ocdbt'):
         asyncio.run(
             ocdbt_utils.merge_ocdbt_per_process_files(
                 test_context.path,
@@ -203,7 +203,7 @@ async def serialize_and_wait():
         logging.info('OCDBT merge complete for %s', test_context.path)
 
     # --- Metadata Validation ---
-    with metrics.time('metadata_validation'):
+    with metrics.measure('metadata_validation'):
       metadata = asyncio.run(handler.metadata([param_info]))[0]
       self._validate_metadata(metadata, sharded_array, param_info.name)
       multihost.sync_global_processes('metadata validation complete')
@@ -215,7 +215,7 @@ async def serialize_and_wait():
         global_shape=sharded_array.shape,
         dtype=sharded_array.dtype,
     )
-    with metrics.time('deserialize'):
+    with metrics.measure('deserialize'):
       restored_array = asyncio.run(
           handler.deserialize([param_info], args=[restore_args])
       )[0]
@@ -224,7 +224,7 @@ async def serialize_and_wait():
     logging.info('Deserialization complete for %s', param_info.name)
 
     # --- Restored Array Validation ---
-    with metrics.time('correctness_check'):
+    with metrics.measure('correctness_check'):
       pytree_utils.assert_pytree_equal(sharded_array, restored_array)
     logging.info('Correctness check passed for %s', param_info.name)
 

checkpoint/orbax/checkpoint/_src/testing/benchmarks/array_handler_benchmark_test.py

@@ -138,10 +138,10 @@ def test_benchmark_returns_test_result_with_basic_metrics(self, options):
     result = self._run_benchmark_workflow_test(options)
 
     self.assertIsInstance(result, benchmarks_core.TestResult)
-    self.assertIn('serialize_time', result.metrics.results)
-    self.assertIn('metadata_validation_time', result.metrics.results)
-    self.assertIn('deserialize_time', result.metrics.results)
-    self.assertIn('correctness_check_time', result.metrics.results)
+    self.assertIn('serialize_time_duration', result.metrics.results)
+    self.assertIn('metadata_validation_time_duration', result.metrics.results)
+    self.assertIn('deserialize_time_duration', result.metrics.results)
+    self.assertIn('correctness_check_time_duration', result.metrics.results)
 
   @parameterized.named_parameters(
       dict(
@@ -156,14 +156,14 @@ def test_benchmark_returns_test_result_with_basic_metrics(self, options):
   def test_benchmark_ocdbt_enabled_calls_merge(self, options):
     result = self._run_benchmark_workflow_test(options)
 
-    self.assertIn('merge_ocdbt_time', result.metrics.results)
+    self.assertIn('merge_ocdbt_time_duration', result.metrics.results)
     self.mock_merge_ocdbt.assert_called_once()
 
   def test_benchmark_ocdbt_disabled_does_not_merge(self):
     options = ArrayHandlerBenchmarkOptions(use_ocdbt=False, use_zarr3=True)
     result = self._run_benchmark_workflow_test(options)
 
-    self.assertNotIn('merge_ocdbt_time', result.metrics.results)
+    self.assertNotIn('merge_ocdbt_time_duration', result.metrics.results)
     self.mock_merge_ocdbt.assert_not_called()
 
   @parameterized.named_parameters(

checkpoint/orbax/checkpoint/_src/testing/benchmarks/checkpoint_manager_benchmark.py

@@ -90,10 +90,10 @@ def test_fn(
     step_saved = -1
     for step in range(options.train_steps):
       logging.info('Saving checkpoint at step %d', step)
-      with metrics.time(f'save_{step}'):
+      with metrics.measure(f'save_{step}'):
         if mngr.save(step, args=composite_args):
           step_saved = step
-      with metrics.time(f'wait_until_finished_{step}'):
+      with metrics.measure(f'wait_until_finished_{step}'):
         mngr.wait_until_finished()
       logging.info('Finished saving checkpoint at step %d', step)
 
@@ -105,12 +105,12 @@ def test_fn(
         f'Expected latest step to be {step_saved}, got {latest_step}'
     )
 
-    with metrics.time(f'restore_{latest_step}'):
+    with metrics.measure(f'restore_{latest_step}'):
       logging.info('Restoring checkpoint at step %d', latest_step)
       restored = mngr.restore(latest_step, args=restore_args)
       logging.info('Finished restoring checkpoint at step %d', latest_step)
 
-    with metrics.time('correctness_check'):
+    with metrics.measure('correctness_check'):
       pytree_utils.assert_pytree_equal(pytree, restored['pytree'])
       assert (
           json_data == restored['json_item']

checkpoint/orbax/checkpoint/_src/testing/benchmarks/checkpoint_manager_benchmark_test.py

@@ -110,10 +110,10 @@ def test_benchmark_test_fn_succeeds(self):
     self.assertIsInstance(result, benchmarks_core.TestResult)
     self.assertContainsSubset(
         {
-            'save_0_time',
-            'wait_until_finished_0_time',
-            'restore_0_time',
-            'correctness_check_time',
+            'save_0_time_duration',
+            'wait_until_finished_0_time_duration',
+            'restore_0_time_duration',
+            'correctness_check_time_duration',
         },
         result.metrics.results.keys(),
     )

checkpoint/orbax/checkpoint/_src/testing/benchmarks/checkpoint_manager_perf_benchmark.py

@@ -110,7 +110,7 @@ def test_fn(
     save_times = []
     total_save_times = []
     for i in range(options.train_steps):
-      with metrics.time(f'save_{i}'):
+      with metrics.measure(f'save_{i}'):
         save_start = time.time()
         mngr.save(
             i,
@@ -121,7 +121,7 @@ def test_fn(
         save_times.append(time.time() - save_start)
         mngr.wait_until_finished()
         total_save_times.append(time.time() - save_start)
-      with metrics.time(f'train_step_{i}'):
+      with metrics.measure(f'train_step_{i}'):
         pytree = self._train_step(pytree)
 
     save_times = np.array(save_times)
@@ -150,7 +150,7 @@ def test_fn(
     )
     context.pytree = self._clear_pytree(context.pytree)
 
-    with metrics.time('restore'):
+    with metrics.measure('restore'):
       mngr.restore(
           mngr.latest_step(),
           args=ocp.args.Composite(

checkpoint/orbax/checkpoint/_src/testing/benchmarks/checkpoint_policy_benchmark.py

@@ -237,7 +237,7 @@ def test_fn(
         ),
     )
     for step in range(options.num_checkpoints):
-      with metrics.time(f'saving step {step}'):
+      with metrics.measure(f'saving step {step}'):
         checkpointer_manager.save(step, args=ocp.args.PyTreeSave(pytree))
     checkpointer_manager.wait_until_finished()
     all_steps = checkpointer_manager.all_steps()

checkpoint/orbax/checkpoint/_src/testing/benchmarks/core/core.py

@@ -143,12 +143,12 @@ def run(self) -> TestResult:
     )
 
     benchmark_metrics = metric_lib.Metrics(name=f"{self.name} Internal")
-    with benchmark_metrics.time("sync_global_processes:benchmark:run"):
+    with benchmark_metrics.measure("sync_global_processes:benchmark:run"):
       multihost.sync_global_processes("benchmark:run")
 
     path = directory_setup.setup_test_directory(self.name, self.output_dir)
 
-    with benchmark_metrics.time(
+    with benchmark_metrics.measure(
         "sync_global_processes:benchmark:setup_test_directory"
     ):
       multihost.sync_global_processes("benchmark:setup_test_directory")
@@ -160,7 +160,9 @@ def run(self) -> TestResult:
     else:
       data = checkpoint_generation.load_checkpoint(self.checkpoint_config.path)
 
-    with benchmark_metrics.time("sync_global_processes:benchmark:setup_pytree"):
+    with benchmark_metrics.measure(
+        "sync_global_processes:benchmark:setup_pytree"
+    ):
       multihost.sync_global_processes("benchmark:setup_pytree")
 
     context = TestContext(