diff --git a/checkpoint/orbax/checkpoint/_src/metadata/array_metadata.py b/checkpoint/orbax/checkpoint/_src/metadata/array_metadata.py index 820b296cb1..7d817603f5 100644 --- a/checkpoint/orbax/checkpoint/_src/metadata/array_metadata.py +++ b/checkpoint/orbax/checkpoint/_src/metadata/array_metadata.py @@ -40,6 +40,8 @@ class ArrayMetadata: ext_metadata: ExtMetadata | None = ( None # to contain any extension metadata for an array. ) + compression_algorithm: str = 'none' + compression_level: int | None = None @dataclasses.dataclass(frozen=True, kw_only=True) diff --git a/checkpoint/orbax/checkpoint/_src/serialization/jax_array_handlers.py b/checkpoint/orbax/checkpoint/_src/serialization/jax_array_handlers.py index 173439ba1a..20e71a057d 100644 --- a/checkpoint/orbax/checkpoint/_src/serialization/jax_array_handlers.py +++ b/checkpoint/orbax/checkpoint/_src/serialization/jax_array_handlers.py @@ -280,6 +280,51 @@ def _record_logical_metrics( ) +def _record_compression_metrics( + direction: types.IoDirection, + logical_bytes: int, + raw_bytes: int, + storage_type: str, + custom_prefix: str = '', + metadatas: Sequence[ts_utils.ArrayMetadata] | None = None, +) -> None: + """Logs and records compression ratio and metrics.""" + if logical_bytes <= 0: + return + ratio = float(raw_bytes) / logical_bytes + algo_str = 'none' + level_str = 'None' + if metadatas is not None: + algo_str, level_str = ts_utils.resolve_compression_settings(metadatas) + logging.info( + '[process=%d] %s ratio (raw/logical): %.3f (%s / %s), algo=%s, level=%s', + multihost.process_index(), + direction.value.capitalize(), + ratio, + humanize.naturalsize(raw_bytes, binary=True), + humanize.naturalsize(logical_bytes, binary=True), + algo_str, + level_str, + ) + jax.monitoring.record_scalar( + f'/jax/orbax/{direction.value}/worker/io/compression_ratio', + ratio, + storage_type=storage_type, + custom_prefix=custom_prefix, + compression_algorithm=algo_str, + compression_level=level_str, + ) + if direction == types.IoDirection.WRITE: + jax.monitoring.record_scalar( + '/jax/orbax/write/worker/io/compressed_gbytes', + raw_bytes / (1024**3), + storage_type=storage_type, + custom_prefix=custom_prefix, + compression_algorithm=algo_str, + compression_level=level_str, + ) + + def _record_raw_metrics( direction: types.IoDirection, logical_bytes: int, @@ -287,27 +332,19 @@ def _record_raw_metrics( storage_type: str, initial_ts_metrics: Sequence[dict[str, Any]] | None = None, custom_prefix: str = '', + metadatas: Sequence[ts_utils.ArrayMetadata] | None = None, ): """Records raw metrics collected from TensorStore.""" if initial_ts_metrics is None: return - try: - final_ts_metrics = ts.experimental_collect_matching_metrics('/tensorstore/') - except Exception: # pylint: disable=broad-except - final_ts_metrics = None - + final_ts_metrics = ts_utils.collect_tensorstore_metrics() if final_ts_metrics is None: return - initial_bytes = ts_utils.get_total_bytes_from_tensorstore( - initial_ts_metrics, direction + raw_bytes = ts_utils.get_tensorstore_raw_bytes_delta( + initial_ts_metrics, final_ts_metrics, direction ) - final_bytes = ts_utils.get_total_bytes_from_tensorstore( - final_ts_metrics, direction - ) - raw_bytes = final_bytes - initial_bytes - if raw_bytes <= 0: return @@ -335,29 +372,14 @@ def _record_raw_metrics( custom_prefix=custom_prefix, ) - if logical_bytes > 0: - ratio = float(raw_bytes) / logical_bytes - logging.info( - '[process=%d] %s ratio (raw/logical): %.3f (%s / %s)', - multihost.process_index(), - direction.value.capitalize(), - ratio, - humanize.naturalsize(raw_bytes, binary=True), - humanize.naturalsize(logical_bytes, binary=True), - ) - jax.monitoring.record_scalar( - f'/jax/orbax/{direction.value}/worker/io/compression_ratio', - ratio, - storage_type=storage_type, - custom_prefix=custom_prefix, - ) - if direction == types.IoDirection.WRITE: - jax.monitoring.record_scalar( - '/jax/orbax/write/worker/io/compressed_gbytes', - raw_bytes / (1024**3), - storage_type=storage_type, - custom_prefix=custom_prefix, - ) + _record_compression_metrics( + direction, + logical_bytes, + raw_bytes, + storage_type, + custom_prefix=custom_prefix, + metadatas=metadatas, + ) def _log_io_metrics( @@ -367,6 +389,7 @@ def _log_io_metrics( parent_dir: epath.Path, initial_ts_metrics: Sequence[dict[str, Any]] | None = None, custom_prefix: str = '', + metadatas: Sequence[ts_utils.ArrayMetadata] | None = None, ): """Logs and records IO telemetry metrics for array serialization/deserialization.""" duration = time.time() - start_time @@ -386,6 +409,7 @@ def _log_io_metrics( storage_type, initial_ts_metrics=initial_ts_metrics, custom_prefix=custom_prefix, + metadatas=metadatas, ) @@ -404,12 +428,7 @@ def _worker_serialize_arrays( ext_metadata: Dict[str, Any], ): """Worker function to serialize arrays.""" - try: - initial_ts_metrics = ts.experimental_collect_matching_metrics( - '/tensorstore/' - ) - except Exception: # pylint: disable=broad-except - initial_ts_metrics = None + initial_ts_metrics = ts_utils.collect_tensorstore_metrics() total_start_time = time.time() rslices_per_array = _get_replica_slices( arrays, @@ -419,7 +438,7 @@ def _worker_serialize_arrays( max_replicas_for_replica_parallel, ) - asyncio_utils.run_sync( + array_metadatas = asyncio_utils.run_sync( _async_serialize_replica_slices( rslices_per_array, infos, @@ -440,6 +459,7 @@ def _worker_serialize_arrays( start_time=total_start_time, parent_dir=infos[0].parent_dir, initial_ts_metrics=initial_ts_metrics, + metadatas=array_metadatas, ) @@ -551,22 +571,20 @@ def _serialize_arrays_batches_without_dispatcher( async def _serialize_without_dispatcher(): if not prioritized and not deprioritized: return - try: - initial_ts_metrics = ts.experimental_collect_matching_metrics( - '/tensorstore/' - ) - except Exception: # pylint: disable=broad-except - initial_ts_metrics = None + initial_ts_metrics = ts_utils.collect_tensorstore_metrics() total_start_time = time.time() logical_bytes = 0 + all_array_metadatas: list[ts_utils.ArrayMetadata] = [] if prioritized_values_on_host: logical_bytes += sum(v.nbytes for v in prioritized_values_on_host) - await async_serialize_replica_slices_batch( + b_metadatas = await async_serialize_replica_slices_batch( prioritized_values_on_host, prioritized_infos, prioritized_args, ) + if b_metadatas: + all_array_metadatas.extend(b_metadatas) _on_batch_callback(prioritized_infos, callback.on_write_end) if deprioritized: assert device_host_max_bytes is not None @@ -586,11 +604,13 @@ async def _serialize_without_dispatcher(): b_arrays_on_host = replica_slices_transfer_arrays_to_host(b_arrays) _on_batch_callback(b_infos, callback.on_transfer_end) logical_bytes += sum(v.nbytes for v in b_arrays_on_host) - await async_serialize_replica_slices_batch( + b_metadatas = await async_serialize_replica_slices_batch( b_arrays_on_host, b_infos, b_args, ) + if b_metadatas: + all_array_metadatas.extend(b_metadatas) _on_batch_callback(b_infos, callback.on_write_end) info_sample = prioritized[0][1] if prioritized else deprioritized[0][1] @@ -600,6 +620,7 @@ async def _serialize_without_dispatcher(): start_time=total_start_time, parent_dir=info_sample.parent_dir, initial_ts_metrics=initial_ts_metrics, + metadatas=all_array_metadatas, ) return future.CommitFutureAwaitingContractedSignals( @@ -783,7 +804,7 @@ async def _async_serialize_replica_slices( enable_replica_parallel_separate_folder: bool, use_replica_parallel: bool, ext_metadata: Dict[str, Any], -) -> None: +) -> Sequence[ts_utils.ArrayMetadata]: """This function contains the logic from ArrayHandler._background_serialize.""" write_coros = [] array_metadatas = [] @@ -884,6 +905,8 @@ async def _async_serialize_replica_slices( ocdbt_transaction.abort() raise + return array_metadatas + def _wrap_random_key_data( array_metadatas: Any, @@ -1014,12 +1037,7 @@ async def _deserialize_arrays( array_metadata_store: array_metadata_store_lib.Store | None, ) -> Sequence[jax.Array]: """Deserializes arrays and applies array_metadata if available.""" - try: - initial_ts_metrics = ts.experimental_collect_matching_metrics( - '/tensorstore/' - ) - except Exception: # pylint: disable=broad-except - initial_ts_metrics = None + initial_ts_metrics = ts_utils.collect_tensorstore_metrics() total_start_time = time.time() async def _async_deserialize( diff --git a/checkpoint/orbax/checkpoint/_src/serialization/tensorstore_utils.py b/checkpoint/orbax/checkpoint/_src/serialization/tensorstore_utils.py index 28f326abac..551dd4b949 100644 --- a/checkpoint/orbax/checkpoint/_src/serialization/tensorstore_utils.py +++ b/checkpoint/orbax/checkpoint/_src/serialization/tensorstore_utils.py @@ -593,56 +593,77 @@ async def open_kv_store( ### Building Zarr array metadata. -def build_zarr_shard_and_chunk_metadata( +def _build_zarr2_metadata( + global_shape: Shape, + chunk_shape: Shape, + use_compression: bool, +) -> JsonSpec: + """Constructs Zarr v2 metadata.""" + # Use default level 1 straight from TensorStore. + compressor = {'id': 'zstd', 'level': 1} if use_compression else None + return { + 'shape': global_shape, + 'chunks': chunk_shape, + 'compressor': compressor, # pyrefly: ignore[bad-assignment] + } + + +def _build_zarr3_metadata( + global_shape: Shape, + chunk_shape: Shape, + use_compression: bool, +) -> JsonSpec: + """Constructs Zarr v3 metadata.""" + codecs: list[JsonSpec] = [{ + 'name': 'sharding_indexed', + 'configuration': { + 'chunk_shape': chunk_shape, + 'codecs': [ + {'name': 'bytes', 'configuration': {'endian': 'little'}}, + ], + 'index_codecs': [ + {'name': 'bytes', 'configuration': {'endian': 'little'}}, + {'name': 'crc32c'}, + ], + 'index_location': 'end', + }, + }] + if use_compression: + # Use default level 3 straight from TensorStore. + codecs[0]['configuration']['codecs'].append( + {'name': 'zstd', 'configuration': {'level': 3}} + ) # pyrefly: ignore[bad-index] + + return { + 'shape': global_shape, + 'chunk_grid': { # pyrefly: ignore[bad-assignment] + 'name': 'regular', + 'configuration': {'chunk_shape': chunk_shape}, + }, + 'codecs': codecs, # pyrefly: ignore[bad-assignment] + } + + +def _build_zarr_shard_and_chunk_metadata( *, global_shape: Shape, shard_shape: Shape, use_compression: bool = True, use_zarr3: bool, chunk_shape: Shape, -) -> JsonSpec: +) -> tuple[JsonSpec, str, int | None]: """Constructs Zarr metadata for TensorStore array write spec.""" - metadata = {'shape': global_shape} - - if not use_zarr3: - # Zarr v2. - metadata['chunks'] = chunk_shape - if use_compression: - metadata['compressor'] = {'id': 'zstd'} # pyrefly: ignore[bad-assignment] - else: - metadata['compressor'] = None # pyrefly: ignore[bad-assignment] + # TODO: b/354139177 - Consider if using write shape equal to shard shape and + # read shape equal to chosen chunk shape would be a better setting. + del shard_shape # Currently unused. + if use_zarr3: + metadata = _build_zarr3_metadata(global_shape, chunk_shape, use_compression) + level = 3 if use_compression else None else: - # Zarr v3. - metadata['chunk_grid'] = { # pyrefly: ignore[bad-assignment] - 'name': 'regular', - 'configuration': { - 'chunk_shape': chunk_shape, - }, - } - # TODO: b/354139177 - Consider if using write shape equal to shard shape and - # read shape equal to chosen chunk shape would be a better setting. - del shard_shape # Currently unused. - metadata['codecs'] = [ # pyrefly: ignore[bad-assignment] - { - 'name': 'sharding_indexed', - 'configuration': { - 'chunk_shape': chunk_shape, - 'codecs': [ - {'name': 'bytes', 'configuration': {'endian': 'little'}}, - ], - 'index_codecs': [ - {'name': 'bytes', 'configuration': {'endian': 'little'}}, - {'name': 'crc32c'}, - ], - 'index_location': 'end', - }, - }, - ] - if use_compression: - # Remove zstd codec if not using compression. - metadata['codecs'][0]['configuration']['codecs'].append({'name': 'zstd'}) # pyrefly: ignore[bad-index] - - return metadata + metadata = _build_zarr2_metadata(global_shape, chunk_shape, use_compression) + level = 1 if use_compression else None + algo = 'zstd' if use_compression else 'none' + return metadata, algo, level def calculate_chunk_byte_size( @@ -851,7 +872,7 @@ def __init__( chunk_shape, ) # Construct Zarr chunk metadata. - tspec['metadata'] = build_zarr_shard_and_chunk_metadata( + tspec['metadata'], algo, level = _build_zarr_shard_and_chunk_metadata( global_shape=global_shape, shard_shape=write_shape, use_compression=use_compression, @@ -869,6 +890,8 @@ def __init__( use_ocdbt=use_ocdbt, use_zarr3=use_zarr3, ext_metadata=ext_metadata, + compression_algorithm=algo, + compression_level=level, ) # Wrap spec into `cast` driver if needed, and keep it in a separate field. self._json_spec = _maybe_add_cast_to_write_spec( @@ -1155,6 +1178,48 @@ def get_total_bytes_from_tensorstore( return total +def get_tensorstore_raw_bytes_delta( + initial_metrics: Sequence[dict[str, Any]] | None, + final_metrics: Sequence[dict[str, Any]] | None, + direction: types.IoDirection = types.IoDirection.WRITE, +) -> int: + """Computes transferred raw bytes delta between two metric snapshots.""" + if initial_metrics is None or final_metrics is None: + return 0 + try: + initial_bytes = get_total_bytes_from_tensorstore(initial_metrics, direction) + final_bytes = get_total_bytes_from_tensorstore(final_metrics, direction) + return max(0, final_bytes - initial_bytes) + except Exception: # pylint: disable=broad-except + logging.exception('Failed to compute TensorStore raw bytes delta.') + return 0 + + +def collect_tensorstore_metrics() -> Sequence[dict[str, Any]] | None: + """Safely collects TensorStore driver metrics.""" + try: + return ts.experimental_collect_matching_metrics('/tensorstore') + except Exception: # pylint: disable=broad-except + return None + + +def resolve_compression_settings( + metadatas: Sequence[ArrayMetadata], +) -> tuple[str, str]: + """Extracts (algo, level) across array metadata.""" + if not metadatas: + return ('none', 'None') + compression_settings = { + (m.compression_algorithm, m.compression_level) for m in metadatas + } + if len(compression_settings) == 1: + algo, level = next(iter(compression_settings)) + return (str(algo), str(level)) + + # this should be rare. + return ('mixed', 'mixed') + + def print_ts_debug_data(key: str | None, infos: Sequence[types.ParamInfo]): """Log Tensorstore related metrics.""" ts_metrics = ts.experimental_collect_matching_metrics('/tensorstore') diff --git a/checkpoint/orbax/checkpoint/_src/serialization/tensorstore_utils_test.py b/checkpoint/orbax/checkpoint/_src/serialization/tensorstore_utils_test.py index a18027b86b..c807e41004 100644 --- a/checkpoint/orbax/checkpoint/_src/serialization/tensorstore_utils_test.py +++ b/checkpoint/orbax/checkpoint/_src/serialization/tensorstore_utils_test.py @@ -125,6 +125,52 @@ def test_get_backend_ocdbt_target_data_file_size( ) +class BuildZarrShardAndChunkMetadataTest(parameterized.TestCase): + + @parameterized.product( + use_zarr3=(True, False), + use_compression=(True, False), + ) + def test_build_zarr_shard_and_chunk_metadata( + self, use_zarr3: bool, use_compression: bool + ): + shape = (10, 6, 32) + chunk_shape = (5, 6, 8) + metadata, algo, level = ts_utils._build_zarr_shard_and_chunk_metadata( + global_shape=shape, + shard_shape=chunk_shape, + use_compression=use_compression, + use_zarr3=use_zarr3, + chunk_shape=chunk_shape, + ) + expected_algo = 'zstd' if use_compression else 'none' + self.assertEqual(algo, expected_algo) + if use_compression: + expected_level = 3 if use_zarr3 else 1 + self.assertEqual(level, expected_level) + else: + self.assertIsNone(level) + self.assertEqual(metadata['shape'], shape) + if not use_zarr3: + self.assertEqual(metadata['chunks'], chunk_shape) + if use_compression: + self.assertEqual(metadata['compressor'], {'id': 'zstd', 'level': 1}) + else: + self.assertIsNone(metadata['compressor']) + else: + self.assertEqual( + metadata['chunk_grid']['configuration']['chunk_shape'], chunk_shape + ) + codecs = metadata['codecs'][0]['configuration']['codecs'] + codec_names = [c['name'] for c in codecs] + if use_compression: + self.assertIn('zstd', codec_names) + zstd_codec = next(c for c in codecs if c['name'] == 'zstd') + self.assertEqual(zstd_codec['configuration']['level'], 3) + else: + self.assertNotIn('zstd', codec_names) + + class BuildArrayTSpecForWriteTest(parameterized.TestCase): def setUp(self): @@ -397,7 +443,7 @@ def test_full_spec_ocdbt_gcs(self, gcs_backend: str): }, 'metadata': { 'chunks': self.write_shape, - 'compressor': {'id': 'zstd'}, + 'compressor': {'id': 'zstd', 'level': 1}, 'shape': self.shape, }, 'recheck_cached_data': False, @@ -654,6 +700,57 @@ def test_casts_to_target_dtype( self.assertEqual(tspec.json['dtype'], 'int32') self.assertEqual(tspec.json['base']['dtype'], 'float32') + @parameterized.product( + use_zarr3=(True, False), + use_ocdbt=(True, False), + ) + def test_compression_uncompressed(self, use_zarr3: bool, use_ocdbt: bool): + tspec = self.array_write_spec_constructor( + directory=self.directory, + relative_array_filename=self.param_name, + use_zarr3=use_zarr3, + use_ocdbt=use_ocdbt, + use_compression=False, + ) + self.assertEqual(tspec.metadata.compression_algorithm, 'none') + self.assertIsNone(tspec.metadata.compression_level) + + @parameterized.product( + use_zarr3=(True, False), + use_ocdbt=(True, False), + ) + def test_compression_default_compressed( + self, use_zarr3: bool, use_ocdbt: bool + ): + tspec = self.array_write_spec_constructor( + directory=self.directory, + relative_array_filename=self.param_name, + use_zarr3=use_zarr3, + use_ocdbt=use_ocdbt, + use_compression=True, + ) + self.assertEqual(tspec.metadata.compression_algorithm, 'zstd') + expected_level = 3 if use_zarr3 else 1 + self.assertEqual(tspec.metadata.compression_level, expected_level) + + @parameterized.product( + use_zarr3=(True, False), + use_ocdbt=(True, False), + ) + def test_compression_with_casting(self, use_zarr3: bool, use_ocdbt: bool): + tspec = self.array_write_spec_constructor( + directory=self.directory, + relative_array_filename=self.param_name, + target_dtype=np.dtype(np.float32), + use_zarr3=use_zarr3, + use_ocdbt=use_ocdbt, + use_compression=True, + ) + self.assertEqual(tspec.json['driver'], 'cast') + self.assertEqual(tspec.metadata.compression_algorithm, 'zstd') + expected_level = 3 if use_zarr3 else 1 + self.assertEqual(tspec.metadata.compression_level, expected_level) + def _get_chunk_shape_from_tspec( self, tspec: ts_utils.JsonSpec, @@ -1415,5 +1512,127 @@ def test_commit_temporary_metadata_mode( self._verify_kvstack_spec(kvstore_tspec['base'], expected_base_path) +class GetTensorStoreRawBytesDeltaTest(parameterized.TestCase): + + def test_none_metrics(self): + self.assertEqual(ts_utils.get_tensorstore_raw_bytes_delta(None, None), 0) + self.assertEqual(ts_utils.get_tensorstore_raw_bytes_delta([], None), 0) + self.assertEqual(ts_utils.get_tensorstore_raw_bytes_delta(None, []), 0) + + def test_delta_calculation(self): + initial = [{ + 'name': '/tensorstore/kvstore/ocdbt/bytes_written', + 'values': [{'value': 100}], + }] + final = [{ + 'name': '/tensorstore/kvstore/ocdbt/bytes_written', + 'values': [{'value': 350}], + }] + delta = ts_utils.get_tensorstore_raw_bytes_delta( + initial, final, serialization_types.IoDirection.WRITE + ) + self.assertEqual(delta, 250) + + def test_negative_delta_returns_zero(self): + initial = [{ + 'name': '/tensorstore/kvstore/ocdbt/bytes_written', + 'values': [{'value': 500}], + }] + final = [{ + 'name': '/tensorstore/kvstore/ocdbt/bytes_written', + 'values': [{'value': 300}], + }] + delta = ts_utils.get_tensorstore_raw_bytes_delta( + initial, final, serialization_types.IoDirection.WRITE + ) + self.assertEqual(delta, 0) + + +class CollectTensorStoreMetricsTest(parameterized.TestCase): + + def test_collect_returns_list_or_none(self): + metrics = ts_utils.collect_tensorstore_metrics() + self.assertTrue(metrics is None or isinstance(metrics, list)) + + +class ResolveCompressionSettingsTest(parameterized.TestCase): + + def test_empty(self): + self.assertEqual( + ts_utils.resolve_compression_settings([]), ('none', 'None') + ) + + def test_single(self): + meta = ts_utils.ArrayMetadata( + param_name='a', + shape=(1,), + dtype=np.dtype('float32'), + write_shape=(1,), + chunk_shape=(1,), + use_ocdbt=True, + use_zarr3=False, + compression_algorithm='zstd', + compression_level=3, + ) + self.assertEqual( + ts_utils.resolve_compression_settings([meta]), ('zstd', '3') + ) + + def test_multiple_identical(self): + meta1 = ts_utils.ArrayMetadata( + param_name='a', + shape=(1,), + dtype=np.dtype('float32'), + write_shape=(1,), + chunk_shape=(1,), + use_ocdbt=True, + use_zarr3=False, + compression_algorithm='zstd', + compression_level=1, + ) + meta2 = ts_utils.ArrayMetadata( + param_name='b', + shape=(2,), + dtype=np.dtype('float32'), + write_shape=(2,), + chunk_shape=(2,), + use_ocdbt=True, + use_zarr3=False, + compression_algorithm='zstd', + compression_level=1, + ) + self.assertEqual( + ts_utils.resolve_compression_settings([meta1, meta2]), ('zstd', '1') + ) + + def test_mixed(self): + meta1 = ts_utils.ArrayMetadata( + param_name='a', + shape=(1,), + dtype=np.dtype('float32'), + write_shape=(1,), + chunk_shape=(1,), + use_ocdbt=True, + use_zarr3=False, + compression_algorithm='zstd', + compression_level=3, + ) + meta2 = ts_utils.ArrayMetadata( + param_name='b', + shape=(2,), + dtype=np.dtype('float32'), + write_shape=(2,), + chunk_shape=(2,), + use_ocdbt=True, + use_zarr3=False, + compression_algorithm='none', + compression_level=None, + ) + self.assertEqual( + ts_utils.resolve_compression_settings([meta1, meta2]), + ('mixed', 'mixed'), + ) + + if __name__ == '__main__': absltest.main()