diff --git a/checkpoint/orbax/checkpoint/experimental/v1/_src/testing/compatibility/checkpointables_metadata_compatibility_test.py b/checkpoint/orbax/checkpoint/experimental/v1/_src/testing/compatibility/checkpointables_metadata_compatibility_test_base.py similarity index 80% rename from checkpoint/orbax/checkpoint/experimental/v1/_src/testing/compatibility/checkpointables_metadata_compatibility_test.py rename to checkpoint/orbax/checkpoint/experimental/v1/_src/testing/compatibility/checkpointables_metadata_compatibility_test_base.py index 41f64811a..54f601de9 100644 --- a/checkpoint/orbax/checkpoint/experimental/v1/_src/testing/compatibility/checkpointables_metadata_compatibility_test.py +++ b/checkpoint/orbax/checkpoint/experimental/v1/_src/testing/compatibility/checkpointables_metadata_compatibility_test_base.py @@ -26,6 +26,7 @@ from orbax.checkpoint.experimental.v1._src.context import options as options_lib from orbax.checkpoint.experimental.v1._src.handlers import registration from orbax.checkpoint.experimental.v1._src.layout import checkpoint_layout as checkpoint_layout_lib +from orbax.checkpoint.experimental.v1._src.synchronization import multihost from orbax.checkpoint.experimental.v1._src.testing.compatibility import test_utils as compatibility_test_utils @@ -36,7 +37,8 @@ _BASE_DIR = os.path.join(os.path.dirname(__file__), 'checkpoints') -class CheckpointablesMetadataCompatibilityTest(parameterized.TestCase): +class CheckpointablesMetadataCompatibilityTestBase(parameterized.TestCase): + """Tests for V1 checkpointables_metadata API against generated Checkpoints.""" def setUp(self) -> None: super().setUp() @@ -94,6 +96,14 @@ def test_checkpointables_metadata_compatibility( is_direct_checkpoint: bool, is_pytree: bool, ) -> None: + """Tests checkpointables_metadata against various checkpoint formats. + + Args: + version: v0 or v1. + metadata_present: Whether the checkpoint has metadata files. + is_direct_checkpoint: Whether the checkpoint is a direct checkpoint. + is_pytree: Whether the checkpoint is a pytree checkpoint. + """ path = compatibility_test_utils.get_checkpoint_path( version, metadata_present, is_direct_checkpoint, is_pytree ) @@ -133,7 +143,11 @@ def test_checkpointables_metadata_compatibility( else: expected = self.expected_checkpointables_metadata - test_utils.assert_tree_equal(self, expected, loaded.metadata) + actual = loaded.metadata + if multihost.is_pathways_backend() or jax.process_count() > 1: + expected = compatibility_test_utils.strip_sharding_metadata(expected) + actual = compatibility_test_utils.strip_sharding_metadata(actual) + test_utils.assert_tree_equal(self, expected, actual) else: with self.assertRaisesRegex(error_type, expected_error_msg): ocp.checkpointables_metadata(path) @@ -153,6 +167,12 @@ def test_checkpointables_metadata_compatibility( def test_checkpointables_metadata_non_critical_corruptions( self, version: str, alteration: str ) -> None: + """Tests checkpointables_metadata against non-critical corruptions. + + Args: + version: The checkpoint version to test against. + alteration: The alteration to apply to the checkpoint. + """ path = self.base_dir.joinpath( f'{version}_checkpoints', 'composite_checkpoint', @@ -162,9 +182,12 @@ def test_checkpointables_metadata_non_critical_corruptions( # Missing sharding metadata results in a pytree identical to expected # values except sharding metadata is None. loaded = ocp.checkpointables_metadata(path) - test_utils.assert_tree_equal( - self, self.expected_checkpointables_metadata, loaded.metadata - ) + expected = self.expected_checkpointables_metadata + actual = loaded.metadata + if multihost.is_pathways_backend() or jax.process_count() > 1: + expected = compatibility_test_utils.strip_sharding_metadata(expected) + actual = compatibility_test_utils.strip_sharding_metadata(actual) + test_utils.assert_tree_equal(self, expected, actual) @parameterized.product( version=['v0', 'v1'], @@ -172,6 +195,11 @@ def test_checkpointables_metadata_non_critical_corruptions( def test_checkpointables_metadata_missing_sharding_corruption( self, version: str ) -> None: + """Tests checkpointables_metadata against missing sharding corruption. + + Args: + version: The checkpoint version to test against. + """ path = self.base_dir.joinpath( f'{version}_checkpoints', 'composite_checkpoint', @@ -193,6 +221,12 @@ def test_checkpointables_metadata_missing_sharding_corruption( def test_checkpointables_metadata_critical_corruptions( self, version: str, alteration: str ) -> None: + """Tests checkpointables_metadata against critical corruptions. + + Args: + version: The checkpoint version to test against. + alteration: The alteration to apply to the checkpoint. + """ path = self.base_dir.joinpath( f'{version}_checkpoints', 'composite_checkpoint', diff --git a/checkpoint/orbax/checkpoint/experimental/v1/_src/testing/compatibility/generate_v0_checkpoint_manager_checkpoints.py b/checkpoint/orbax/checkpoint/experimental/v1/_src/testing/compatibility/generate_v0_checkpoint_manager_checkpoints.py new file mode 100644 index 000000000..9f5d9821e --- /dev/null +++ b/checkpoint/orbax/checkpoint/experimental/v1/_src/testing/compatibility/generate_v0_checkpoint_manager_checkpoints.py @@ -0,0 +1,145 @@ +# Copyright 2026 The Orbax Authors. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Generates V0 CheckpointManager checkpoints for compatibility testing. + +The checkpoints generated by this script are checked into the repository +statically to ensure long-term backward compatibility against runtime changes. +""" +import itertools +import os +from typing import Any + +from absl import app +from absl import flags +from etils import epath +import jax +import jax.numpy as jnp +from orbax.checkpoint import args +from orbax.checkpoint import checkpoint_manager +from orbax.checkpoint import handlers +from orbax.checkpoint import test_utils +from orbax.checkpoint.experimental.v1._src.testing import handler_utils + +FLAGS = flags.FLAGS + + +def _get_base_dir(): + if 'BUILD_WORKING_DIRECTORY' in os.environ: + return os.path.join( + os.environ['BUILD_WORKING_DIRECTORY'], + 'orbax/checkpoint/experimental/v1/_src/testing/compatibility/managed_checkpoints', + ) + return os.path.join( + os.path.dirname(__file__), + 'managed_checkpoints', + ) + + +_BASE_DIR = flags.DEFINE_string( + 'base_dir', + _get_base_dir(), + 'Base directory to save checkpoints.', +) +_OVERWRITE = flags.DEFINE_bool( + 'overwrite', + False, + 'Overwrite existing checkpoints.', +) + + +def create_pytree() -> dict[str, Any]: + return { + 'a': jax.device_put(jnp.array([0, 1, 2, 3, 4, 5, 6, 7], dtype=jnp.int32)), + 'b': {'c': jax.device_put(jnp.array([1, 2, 3], dtype=jnp.int32))}, + } + + +def generate_v0_checkpoint_manager_checkpoint( + path: epath.Path, + has_metrics: bool = True, + has_custom_metadata: bool = True, +) -> None: + """Saves a composite manager checkpoint using V0 CheckpointManager.""" + if _OVERWRITE.value: + path.rmtree(missing_ok=True) + path.mkdir(parents=True, exist_ok=True) + + pytree = create_pytree() + json_object = {'metadata': 'json_data'} + baz = handler_utils.Baz(123, 'hi') + registry = handlers.create_default_handler_registry( + state=handlers.PyTreeCheckpointHandler(), + metadata=handlers.JsonCheckpointHandler(), + baz=handler_utils.DataclassCheckpointHandler(), + ) + options = checkpoint_manager.CheckpointManagerOptions( + create=True, + enable_async_checkpointing=False, + prevent_write_metrics=not has_metrics, + best_fn=lambda metrics: metrics['loss'] if has_metrics else None, + ) + + manager = checkpoint_manager.CheckpointManager( + path, + options=options, + handler_registry=registry, + metadata={'foo': 'bar'} if has_custom_metadata else None, + ) + with manager: + manager.save( + 0, + args=args.Composite( + state=args.PyTreeSave(pytree), + metadata=args.JsonSave(json_object), + baz=handler_utils.DataclassSaveArgs(baz), + ), + metrics={'loss': 0.5} if has_metrics else None, + custom_metadata={'custom': 'meta'} if has_custom_metadata else None, + ) + manager.wait_until_finished() + + +def main(argv): + del argv + epath.Path(_BASE_DIR.value).mkdir(parents=True, exist_ok=True) + + test_utils.set_tensorstore_driver_for_test() + + print('Generating V0 CheckpointManager Checkpoints...') + + for has_metrics, has_custom_metadata in itertools.product( + [True, False], repeat=2 + ): + metrics_str = 'has_metrics' if has_metrics else 'no_metrics' + meta_str = ( + 'has_custom_metadata' if has_custom_metadata else 'no_custom_metadata' + ) + path = ( + epath.Path(_BASE_DIR.value) + / 'v0_managed' + / metrics_str + / meta_str + ) + generate_v0_checkpoint_manager_checkpoint( + path, + has_metrics=has_metrics, + has_custom_metadata=has_custom_metadata, + ) + + print(f'V0 CheckpointManager Checkpoints generated at {_BASE_DIR.value}') + + +if __name__ == '__main__': + app.run(main) diff --git a/checkpoint/orbax/checkpoint/experimental/v1/_src/testing/compatibility/generate_v1_checkpointer_checkpoints.py b/checkpoint/orbax/checkpoint/experimental/v1/_src/testing/compatibility/generate_v1_checkpointer_checkpoints.py new file mode 100644 index 000000000..68a515b21 --- /dev/null +++ b/checkpoint/orbax/checkpoint/experimental/v1/_src/testing/compatibility/generate_v1_checkpointer_checkpoints.py @@ -0,0 +1,140 @@ +# Copyright 2026 The Orbax Authors. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Generates V1 Checkpointer checkpoints for compatibility testing. + +The checkpoints generated by this script are checked into the repository +statically to ensure long-term backward compatibility against runtime changes. +""" +import itertools +import os +from typing import Any + +from absl import app +from absl import flags +from etils import epath +import jax +import jax.numpy as jnp +from orbax.checkpoint import test_utils +import orbax.checkpoint.experimental.v1 as ocp +from orbax.checkpoint.experimental.v1._src.testing import handler_utils + +FLAGS = flags.FLAGS + + +def _get_base_dir(): + if 'BUILD_WORKING_DIRECTORY' in os.environ: + return os.path.join( + os.environ['BUILD_WORKING_DIRECTORY'], + 'orbax/checkpoint/experimental/v1/_src/testing/compatibility/managed_checkpoints', + ) + return os.path.join( + os.path.dirname(__file__), + 'managed_checkpoints', + ) + + +_BASE_DIR = flags.DEFINE_string( + 'base_dir', + _get_base_dir(), + 'Base directory to save checkpoints.', +) +_OVERWRITE = flags.DEFINE_bool( + 'overwrite', + False, + 'Overwrite existing checkpoints.', +) + + +def create_pytree() -> dict[str, Any]: + return { + 'a': jax.device_put(jnp.array([0, 1, 2, 3, 4, 5, 6, 7], dtype=jnp.int32)), + 'b': {'c': jax.device_put(jnp.array([1, 2, 3], dtype=jnp.int32))}, + } + + +def generate_v1_checkpointer_checkpoint( + path: epath.Path, + has_metrics: bool = True, + has_custom_metadata: bool = True, +) -> None: + """Saves a composite checkpoint using V1 Checkpointer.""" + if _OVERWRITE.value: + path.rmtree(missing_ok=True) + path.mkdir(parents=True, exist_ok=True) + + pytree = create_pytree() + json_object = {'metadata': 'json_data'} + baz = handler_utils.Baz(123, 'hi') + + registry = ocp.handlers.local_registry() + registry.add(ocp.handlers.PyTreeHandler, checkpointable_name='state') + registry.add(ocp.handlers.JsonHandler, checkpointable_name='metadata') + registry.add(handler_utils.BazHandler, checkpointable_name='baz') + + checkpointables = { + 'state': pytree, + 'metadata': json_object, + 'baz': baz, + } + + with ocp.Context( + checkpointables_options=ocp.options.CheckpointablesOptions( + registry=registry + ) + ): + checkpointer = ocp.training.Checkpointer( + path, + custom_metadata={'foo': 'bar'} if has_custom_metadata else None, + ) + checkpointer.save_checkpointables( + 0, + checkpointables, + metrics={'loss': 0.5} if has_metrics else None, + custom_metadata={'custom': 'meta'} if has_custom_metadata else None, + ) + + +def main(argv): + del argv + epath.Path(_BASE_DIR.value).mkdir(parents=True, exist_ok=True) + + test_utils.set_tensorstore_driver_for_test() + + print('Generating V1 Checkpointer Checkpoints...') + + for has_metrics, has_custom_metadata in itertools.product( + [True, False], repeat=2 + ): + metrics_str = 'has_metrics' if has_metrics else 'no_metrics' + meta_str = ( + 'has_custom_metadata' if has_custom_metadata else 'no_custom_metadata' + ) + path = ( + epath.Path(_BASE_DIR.value) + / 'v1_managed' + / metrics_str + / meta_str + ) + generate_v1_checkpointer_checkpoint( + path, + has_metrics=has_metrics, + has_custom_metadata=has_custom_metadata, + ) + + print(f'V1 Checkpointer Checkpoints generated at {_BASE_DIR.value}') + + +if __name__ == '__main__': + app.run(main) diff --git a/checkpoint/orbax/checkpoint/experimental/v1/_src/testing/compatibility/load_checkpointables_compatibility_test.py b/checkpoint/orbax/checkpoint/experimental/v1/_src/testing/compatibility/load_checkpointables_compatibility_test_base.py similarity index 83% rename from checkpoint/orbax/checkpoint/experimental/v1/_src/testing/compatibility/load_checkpointables_compatibility_test.py rename to checkpoint/orbax/checkpoint/experimental/v1/_src/testing/compatibility/load_checkpointables_compatibility_test_base.py index 2868df6c8..ffe275ea0 100644 --- a/checkpoint/orbax/checkpoint/experimental/v1/_src/testing/compatibility/load_checkpointables_compatibility_test.py +++ b/checkpoint/orbax/checkpoint/experimental/v1/_src/testing/compatibility/load_checkpointables_compatibility_test_base.py @@ -20,12 +20,14 @@ from etils import epath import jax import jax.numpy as jnp +import numpy as np from orbax.checkpoint import test_utils from orbax.checkpoint._src.sharding_utils import make_single_device_sharding import orbax.checkpoint.experimental.v1 as ocp from orbax.checkpoint.experimental.v1._src.context import options as options_lib from orbax.checkpoint.experimental.v1._src.handlers import registration from orbax.checkpoint.experimental.v1._src.layout import checkpoint_layout as checkpoint_layout_lib +from orbax.checkpoint.experimental.v1._src.synchronization import multihost from orbax.checkpoint.experimental.v1._src.testing.compatibility import test_utils as compatibility_test_utils @@ -36,7 +38,8 @@ _BASE_DIR = os.path.join(os.path.dirname(__file__), 'checkpoints') -class LoadCheckpointablesCompatibilityTest(parameterized.TestCase): +class LoadCheckpointablesCompatibilityTestBase(parameterized.TestCase): + """Tests for V1 load_checkpointables API against generated Checkpoints.""" def setUp(self) -> None: super().setUp() @@ -46,6 +49,12 @@ def setUp(self) -> None: 'b': {'c': jnp.array([1, 2, 3], dtype=jnp.int32)}, } sharding = make_single_device_sharding(jax.devices()[0]) + if multihost.is_pathways_backend() or jax.process_count() > 1: + self.expected_state = compatibility_test_utils.replicate_on_mesh( + self.expected_state + ) + mesh = jax.sharding.Mesh(np.asarray(jax.devices()), ('devices',)) + sharding = jax.sharding.NamedSharding(mesh, jax.sharding.PartitionSpec()) self.abstract_state = jax.tree.map( lambda x: jax.ShapeDtypeStruct(x.shape, x.dtype, sharding=sharding), self.expected_state @@ -177,12 +186,32 @@ def test_load_checkpointables_compatibility( has_pytree_metadata: bool, handler_registered: bool, ) -> None: + """Tests load_checkpointables compatibility with v0 and v1 checkpoints. + + Args: + version: The version of the checkpoint to load. + checkpointable_names_provided: Whether to provide checkpointable_names to + load_checkpointables. + abstract_checkpointables_provided: Whether to provide + abstract_checkpointables to ocp.load_checkpointables. + names_registered: Whether to register all checkpointable names to a + handler. + metadata_present: Whether the checkpoint has metadata files. + is_direct_checkpoint: Whether the checkpoint is a direct checkpoint. + has_pytree_metadata: Whether the checkpoint has pytree metadata files. + handler_registered: Whether to register all handlers. + """ path = compatibility_test_utils.get_checkpoint_path( version, metadata_present, is_direct_checkpoint, has_pytree_metadata ) if path is None or not path.exists(): self.skipTest('Checkpoint for combination does not exist.') + if ( + multihost.is_pathways_backend() or jax.process_count() > 1 + ) and not abstract_checkpointables_provided: + self.skipTest('Sharding metadata not matching in Pathways/Multiprocess.') + if not checkpointable_names_provided and abstract_checkpointables_provided: self.skipTest( 'Cannot provide abstract_checkpointables without' @@ -252,6 +281,12 @@ def test_load_checkpointables_compatibility( def test_load_checkpointables_non_critical_corruptions( self, version: str, alteration: str ) -> None: + """Tests load_checkpointables with non-critical corruptions. + + Args: + version: The checkpoint version to test against. + alteration: The alteration to apply to the checkpoint. + """ path = self.base_dir.joinpath( f'{version}_checkpoints', 'composite_checkpoint', @@ -261,7 +296,9 @@ def test_load_checkpointables_non_critical_corruptions( loaded = ocp.load_checkpointables( path, abstract_checkpointables=self.abstract_checkpointables ) - test_utils.assert_tree_equal(self, loaded, self.expected_checkpointables) + test_utils.assert_tree_equal( + self, loaded, self.expected_checkpointables + ) @parameterized.product( version=['v0', 'v1'], @@ -273,6 +310,12 @@ def test_load_checkpointables_non_critical_corruptions( def test_load_checkpointables_critical_corruptions( self, version: str, alteration: str ) -> None: + """Tests load_checkpointables with critical corruptions. + + Args: + version: The checkpoint version to test against. + alteration: The alteration to apply to the checkpoint. + """ path = self.base_dir.joinpath( f'{version}_checkpoints', 'composite_checkpoint', diff --git a/checkpoint/orbax/checkpoint/experimental/v1/_src/testing/compatibility/load_pytree_compatibility_test.py b/checkpoint/orbax/checkpoint/experimental/v1/_src/testing/compatibility/load_pytree_compatibility_test_base.py similarity index 83% rename from checkpoint/orbax/checkpoint/experimental/v1/_src/testing/compatibility/load_pytree_compatibility_test.py rename to checkpoint/orbax/checkpoint/experimental/v1/_src/testing/compatibility/load_pytree_compatibility_test_base.py index 5d82a33f5..c9a63def2 100644 --- a/checkpoint/orbax/checkpoint/experimental/v1/_src/testing/compatibility/load_pytree_compatibility_test.py +++ b/checkpoint/orbax/checkpoint/experimental/v1/_src/testing/compatibility/load_pytree_compatibility_test_base.py @@ -13,6 +13,7 @@ # limitations under the License. """Tests for V1 load_pytree API against generated V0 and V1 Checkpoints.""" + import os from typing import Tuple, Type @@ -21,12 +22,14 @@ from etils import epath import jax import jax.numpy as jnp +import numpy as np from orbax.checkpoint import test_utils from orbax.checkpoint._src.sharding_utils import make_single_device_sharding import orbax.checkpoint.experimental.v1 as ocp from orbax.checkpoint.experimental.v1._src.context import options as options_lib from orbax.checkpoint.experimental.v1._src.handlers import registration from orbax.checkpoint.experimental.v1._src.layout import checkpoint_layout as checkpoint_layout_lib +from orbax.checkpoint.experimental.v1._src.synchronization import multihost from orbax.checkpoint.experimental.v1._src.testing.compatibility import test_utils as compatibility_test_utils @@ -37,7 +40,8 @@ _BASE_DIR = os.path.join(os.path.dirname(__file__), 'checkpoints') -class LoadPytreeCompatibilityTest(parameterized.TestCase): +class LoadPytreeCompatibilityTestBase(parameterized.TestCase): + """Tests for V1 load_pytree API against generated Checkpoints.""" def setUp(self) -> None: super().setUp() @@ -47,6 +51,12 @@ def setUp(self) -> None: 'b': {'c': jnp.array([1, 2, 3], dtype=jnp.int32)}, } sharding = make_single_device_sharding(jax.devices()[0]) + if multihost.is_pathways_backend() or jax.process_count() > 1: + self.expected_state = compatibility_test_utils.replicate_on_mesh( + self.expected_state + ) + mesh = jax.sharding.Mesh(np.asarray(jax.devices()), ('devices',)) + sharding = jax.sharding.NamedSharding(mesh, jax.sharding.PartitionSpec()) self.abstract_state = jax.tree.map( lambda x: jax.ShapeDtypeStruct(x.shape, x.dtype, sharding=sharding), self.expected_state @@ -194,12 +204,33 @@ def test_load_pytree_compatibility( handler_registered: bool, pytree_registered: bool, ) -> None: + """Tests load_pytree against various checkpoint configurations. + + Args: + version: The checkpoint version to test against. + checkpointable_name: The name of the checkpointable to load. + abstract_pytree_provided: Whether an abstract pytree is provided to + ocp.load_pytree. + name_registered: Whether a handler is registered for the + checkpointable_name. + metadata_present: Whether the checkpoint has metadata. + is_direct_checkpoint: Whether the checkpoint is a direct checkpoint. + is_pytree: Whether the checkpoint is a pytree checkpoint. + handler_registered: Whether a handler is registered for the handler + typestr derived from checkpoint metadata. + pytree_registered: Whether a handler is registered for the 'pytree' scope. + """ path = compatibility_test_utils.get_checkpoint_path( version, metadata_present, is_direct_checkpoint, is_pytree ) if path is None or not path.exists(): self.skipTest('Checkpoint for combination does not exist.') + if ( + multihost.is_pathways_backend() or jax.process_count() > 1 + ) and not abstract_pytree_provided: + self.skipTest('Sharding metadata not matching in Pathways/Multiprocess.') + registry = self.setup_registry( path, checkpointable_name, @@ -262,7 +293,12 @@ def test_load_pytree_compatibility( def test_load_pytree_non_critical_corruptions( self, version: str, alteration: str ) -> None: + """Tests load_pytree against checkpoints with non-critical corruptions. + Args: + version: The version of the checkpoint to load. + alteration: The non-critical corruption alteration to apply. + """ path = self.base_dir.joinpath( f'{version}_checkpoints', 'composite_checkpoint', @@ -284,6 +320,12 @@ def test_load_pytree_non_critical_corruptions( def test_load_pytree_critical_corruptions( self, version: str, alteration: str ) -> None: + """Tests load_pytree against checkpoints with critical corruptions. + + Args: + version: The version of the checkpoint to load. + alteration: The critical corruption alteration to apply. + """ path = self.base_dir.joinpath( f'{version}_checkpoints', 'composite_checkpoint', @@ -303,6 +345,11 @@ def test_load_pytree_critical_corruptions( version=['v0', 'v1'], ) def test_load_incorrect_path(self, version: str) -> None: + """Tests load_pytree against checkpoints with incorrect paths. + + Args: + version: The version of the checkpoint to test against. + """ checkpoint_path = ( self.base_dir / f'{version}_checkpoints' diff --git a/checkpoint/orbax/checkpoint/experimental/v1/_src/testing/compatibility/managed_checkpoints/v0_managed/has_metrics/has_custom_metadata/0/_CHECKPOINT_METADATA b/checkpoint/orbax/checkpoint/experimental/v1/_src/testing/compatibility/managed_checkpoints/v0_managed/has_metrics/has_custom_metadata/0/_CHECKPOINT_METADATA new file mode 100644 index 000000000..42e25ca6b --- /dev/null +++ b/checkpoint/orbax/checkpoint/experimental/v1/_src/testing/compatibility/managed_checkpoints/v0_managed/has_metrics/has_custom_metadata/0/_CHECKPOINT_METADATA @@ -0,0 +1 @@ +{"item_handlers": {"baz": "orbax.checkpoint.experimental.v1._src.testing.handler_utils.DataclassCheckpointHandler", "metadata": "orbax.checkpoint._src.handlers.json_checkpoint_handler.JsonCheckpointHandler", "metrics": "orbax.checkpoint._src.handlers.json_checkpoint_handler.JsonCheckpointHandler", "state": "orbax.checkpoint._src.handlers.pytree_checkpoint_handler.PyTreeCheckpointHandler"}, "metrics": {}, "performance_metrics": {}, "init_timestamp_nsecs": 1778139888892258082, "commit_timestamp_nsecs": 1778139890858878738, "custom_metadata": {"custom": "meta"}} \ No newline at end of file diff --git a/checkpoint/orbax/checkpoint/experimental/v1/_src/testing/compatibility/managed_checkpoints/v0_managed/has_metrics/has_custom_metadata/0/baz/foo.txt b/checkpoint/orbax/checkpoint/experimental/v1/_src/testing/compatibility/managed_checkpoints/v0_managed/has_metrics/has_custom_metadata/0/baz/foo.txt new file mode 100644 index 000000000..c9b962c5a --- /dev/null +++ b/checkpoint/orbax/checkpoint/experimental/v1/_src/testing/compatibility/managed_checkpoints/v0_managed/has_metrics/has_custom_metadata/0/baz/foo.txt @@ -0,0 +1 @@ +{"int_val": 123, "str_val": "hi"} \ No newline at end of file diff --git a/checkpoint/orbax/checkpoint/experimental/v1/_src/testing/compatibility/managed_checkpoints/v0_managed/has_metrics/has_custom_metadata/0/metadata/metadata b/checkpoint/orbax/checkpoint/experimental/v1/_src/testing/compatibility/managed_checkpoints/v0_managed/has_metrics/has_custom_metadata/0/metadata/metadata new file mode 100644 index 000000000..7ac5d3905 --- /dev/null +++ b/checkpoint/orbax/checkpoint/experimental/v1/_src/testing/compatibility/managed_checkpoints/v0_managed/has_metrics/has_custom_metadata/0/metadata/metadata @@ -0,0 +1 @@ +{"metadata": "json_data"} \ No newline at end of file diff --git a/checkpoint/orbax/checkpoint/experimental/v1/_src/testing/compatibility/managed_checkpoints/v0_managed/has_metrics/has_custom_metadata/0/metrics/metrics b/checkpoint/orbax/checkpoint/experimental/v1/_src/testing/compatibility/managed_checkpoints/v0_managed/has_metrics/has_custom_metadata/0/metrics/metrics new file mode 100644 index 000000000..07c48241d --- /dev/null +++ b/checkpoint/orbax/checkpoint/experimental/v1/_src/testing/compatibility/managed_checkpoints/v0_managed/has_metrics/has_custom_metadata/0/metrics/metrics @@ -0,0 +1 @@ +{"loss": 0.5} \ No newline at end of file diff --git a/checkpoint/orbax/checkpoint/experimental/v1/_src/testing/compatibility/managed_checkpoints/v0_managed/has_metrics/has_custom_metadata/0/state/_METADATA b/checkpoint/orbax/checkpoint/experimental/v1/_src/testing/compatibility/managed_checkpoints/v0_managed/has_metrics/has_custom_metadata/0/state/_METADATA new file mode 100644 index 000000000..d18275eb9 --- /dev/null +++ b/checkpoint/orbax/checkpoint/experimental/v1/_src/testing/compatibility/managed_checkpoints/v0_managed/has_metrics/has_custom_metadata/0/state/_METADATA @@ -0,0 +1 @@ +{"tree_metadata": {"('a',)": {"key_metadata": [{"key": "a", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8]}}, "('b', 'c')": {"key_metadata": [{"key": "b", "key_type": 2}, {"key": "c", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3]}}}, "use_ocdbt": true, "use_zarr3": false, "store_array_data_equal_to_fill_value": true, "custom_metadata": null} \ No newline at end of file diff --git a/checkpoint/orbax/checkpoint/experimental/v1/_src/testing/compatibility/managed_checkpoints/v0_managed/has_metrics/has_custom_metadata/0/state/_sharding b/checkpoint/orbax/checkpoint/experimental/v1/_src/testing/compatibility/managed_checkpoints/v0_managed/has_metrics/has_custom_metadata/0/state/_sharding new file mode 100644 index 000000000..57439c751 --- /dev/null +++ b/checkpoint/orbax/checkpoint/experimental/v1/_src/testing/compatibility/managed_checkpoints/v0_managed/has_metrics/has_custom_metadata/0/state/_sharding @@ -0,0 +1 @@ +{"YQ==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cpu:0\"}","Yi5j":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cpu:0\"}"} \ No newline at end of file diff --git a/checkpoint/orbax/checkpoint/experimental/v1/_src/testing/compatibility/managed_checkpoints/v0_managed/has_metrics/has_custom_metadata/0/state/array_metadatas/process_0 b/checkpoint/orbax/checkpoint/experimental/v1/_src/testing/compatibility/managed_checkpoints/v0_managed/has_metrics/has_custom_metadata/0/state/array_metadatas/process_0 new file mode 100644 index 000000000..9e1980540 --- /dev/null +++ b/checkpoint/orbax/checkpoint/experimental/v1/_src/testing/compatibility/managed_checkpoints/v0_managed/has_metrics/has_custom_metadata/0/state/array_metadatas/process_0 @@ -0,0 +1 @@ +{"array_metadatas": [{"array_metadata": {"param_name": "a", "write_shape": [8], "chunk_shape": [8], "ext_metadata": null}}, {"array_metadata": {"param_name": "b.c", "write_shape": [3], "chunk_shape": [3], "ext_metadata": null}}]} \ No newline at end of file diff --git a/checkpoint/orbax/checkpoint/experimental/v1/_src/testing/compatibility/managed_checkpoints/v0_managed/has_metrics/has_custom_metadata/0/state/d/ea65b6b4d0fbd9cd118027d3440895d5 b/checkpoint/orbax/checkpoint/experimental/v1/_src/testing/compatibility/managed_checkpoints/v0_managed/has_metrics/has_custom_metadata/0/state/d/ea65b6b4d0fbd9cd118027d3440895d5 new file mode 100644 index 000000000..57b9216da Binary files /dev/null and b/checkpoint/orbax/checkpoint/experimental/v1/_src/testing/compatibility/managed_checkpoints/v0_managed/has_metrics/has_custom_metadata/0/state/d/ea65b6b4d0fbd9cd118027d3440895d5 differ diff --git a/checkpoint/orbax/checkpoint/experimental/v1/_src/testing/compatibility/managed_checkpoints/v0_managed/has_metrics/has_custom_metadata/0/state/descriptor/descriptor.pbtxt b/checkpoint/orbax/checkpoint/experimental/v1/_src/testing/compatibility/managed_checkpoints/v0_managed/has_metrics/has_custom_metadata/0/state/descriptor/descriptor.pbtxt new file mode 100644 index 000000000..107a8afb2 --- /dev/null +++ b/checkpoint/orbax/checkpoint/experimental/v1/_src/testing/compatibility/managed_checkpoints/v0_managed/has_metrics/has_custom_metadata/0/state/descriptor/descriptor.pbtxt @@ -0,0 +1,17 @@ +uuid: "566a6305-6aa3-43ee-bfb0-b020084b7a6f" +build_data { + timestamp { + seconds: 1778139795 + } + user: "angelmau" + hostname: "otbfg4.prod.google.com" + path: "/google/src/cloud/buildrabbit-username/buildrabbit-client/" + target: "//orbax/checkpoint/experimental/v1/_src/testing/compatibility:generate_v0_checkpoint_manager_checkpoints" + invocation_id: "a580b822-1e3a-4824-b47a-9bd884c066c7" + changelist: 911781978 + baseline_cl: 911781978 + workspace_id: "angelmau/345866" + citc_snapshot: 679 + client_status: BUILD_CLIENT_STATUS_UNSPECIFIED + verifiable: NON_VERIFIABLE +} diff --git a/checkpoint/orbax/checkpoint/experimental/v1/_src/testing/compatibility/managed_checkpoints/v0_managed/has_metrics/has_custom_metadata/0/state/descriptor/uuid-566a6305-6aa3-43ee-bfb0-b020084b7a6f b/checkpoint/orbax/checkpoint/experimental/v1/_src/testing/compatibility/managed_checkpoints/v0_managed/has_metrics/has_custom_metadata/0/state/descriptor/uuid-566a6305-6aa3-43ee-bfb0-b020084b7a6f new file mode 100644 index 000000000..788a1cc27 --- /dev/null +++ b/checkpoint/orbax/checkpoint/experimental/v1/_src/testing/compatibility/managed_checkpoints/v0_managed/has_metrics/has_custom_metadata/0/state/descriptor/uuid-566a6305-6aa3-43ee-bfb0-b020084b7a6f @@ -0,0 +1 @@ +uuid: "566a6305-6aa3-43ee-bfb0-b020084b7a6f" diff --git a/checkpoint/orbax/checkpoint/experimental/v1/_src/testing/compatibility/managed_checkpoints/v0_managed/has_metrics/has_custom_metadata/0/state/manifest.ocdbt b/checkpoint/orbax/checkpoint/experimental/v1/_src/testing/compatibility/managed_checkpoints/v0_managed/has_metrics/has_custom_metadata/0/state/manifest.ocdbt new file mode 100644 index 000000000..f9d0cf890 Binary files /dev/null and b/checkpoint/orbax/checkpoint/experimental/v1/_src/testing/compatibility/managed_checkpoints/v0_managed/has_metrics/has_custom_metadata/0/state/manifest.ocdbt differ diff --git a/checkpoint/orbax/checkpoint/experimental/v1/_src/testing/compatibility/managed_checkpoints/v0_managed/has_metrics/has_custom_metadata/0/state/ocdbt.process_0/d/072a9af32a8d72943571b7cd43cb57ea b/checkpoint/orbax/checkpoint/experimental/v1/_src/testing/compatibility/managed_checkpoints/v0_managed/has_metrics/has_custom_metadata/0/state/ocdbt.process_0/d/072a9af32a8d72943571b7cd43cb57ea new file mode 100644 index 000000000..8d6d8647a Binary files /dev/null and b/checkpoint/orbax/checkpoint/experimental/v1/_src/testing/compatibility/managed_checkpoints/v0_managed/has_metrics/has_custom_metadata/0/state/ocdbt.process_0/d/072a9af32a8d72943571b7cd43cb57ea differ diff --git a/checkpoint/orbax/checkpoint/experimental/v1/_src/testing/compatibility/managed_checkpoints/v0_managed/has_metrics/has_custom_metadata/0/state/ocdbt.process_0/d/52008c008bd643144f62fe4ee05ddaea b/checkpoint/orbax/checkpoint/experimental/v1/_src/testing/compatibility/managed_checkpoints/v0_managed/has_metrics/has_custom_metadata/0/state/ocdbt.process_0/d/52008c008bd643144f62fe4ee05ddaea new file mode 100644 index 000000000..57b9216da Binary files /dev/null and b/checkpoint/orbax/checkpoint/experimental/v1/_src/testing/compatibility/managed_checkpoints/v0_managed/has_metrics/has_custom_metadata/0/state/ocdbt.process_0/d/52008c008bd643144f62fe4ee05ddaea differ diff --git a/checkpoint/orbax/checkpoint/experimental/v1/_src/testing/compatibility/managed_checkpoints/v0_managed/has_metrics/has_custom_metadata/0/state/ocdbt.process_0/d/e85e76a71fa324686df57072d434bd1f b/checkpoint/orbax/checkpoint/experimental/v1/_src/testing/compatibility/managed_checkpoints/v0_managed/has_metrics/has_custom_metadata/0/state/ocdbt.process_0/d/e85e76a71fa324686df57072d434bd1f new file mode 100644 index 000000000..34ec758d3 Binary files /dev/null and b/checkpoint/orbax/checkpoint/experimental/v1/_src/testing/compatibility/managed_checkpoints/v0_managed/has_metrics/has_custom_metadata/0/state/ocdbt.process_0/d/e85e76a71fa324686df57072d434bd1f differ diff --git a/checkpoint/orbax/checkpoint/experimental/v1/_src/testing/compatibility/managed_checkpoints/v0_managed/has_metrics/has_custom_metadata/0/state/ocdbt.process_0/manifest.ocdbt b/checkpoint/orbax/checkpoint/experimental/v1/_src/testing/compatibility/managed_checkpoints/v0_managed/has_metrics/has_custom_metadata/0/state/ocdbt.process_0/manifest.ocdbt new file mode 100644 index 000000000..4de7da1e3 Binary files /dev/null and b/checkpoint/orbax/checkpoint/experimental/v1/_src/testing/compatibility/managed_checkpoints/v0_managed/has_metrics/has_custom_metadata/0/state/ocdbt.process_0/manifest.ocdbt differ diff --git a/checkpoint/orbax/checkpoint/experimental/v1/_src/testing/compatibility/managed_checkpoints/v0_managed/has_metrics/has_custom_metadata/metadata/_ROOT_METADATA b/checkpoint/orbax/checkpoint/experimental/v1/_src/testing/compatibility/managed_checkpoints/v0_managed/has_metrics/has_custom_metadata/metadata/_ROOT_METADATA new file mode 100644 index 000000000..0b4608849 --- /dev/null +++ b/checkpoint/orbax/checkpoint/experimental/v1/_src/testing/compatibility/managed_checkpoints/v0_managed/has_metrics/has_custom_metadata/metadata/_ROOT_METADATA @@ -0,0 +1 @@ +{"internal_metadata": {}, "custom_metadata": {"foo": "bar"}} \ No newline at end of file diff --git a/checkpoint/orbax/checkpoint/experimental/v1/_src/testing/compatibility/managed_checkpoints/v0_managed/has_metrics/no_custom_metadata/0/_CHECKPOINT_METADATA b/checkpoint/orbax/checkpoint/experimental/v1/_src/testing/compatibility/managed_checkpoints/v0_managed/has_metrics/no_custom_metadata/0/_CHECKPOINT_METADATA new file mode 100644 index 000000000..5ef3142ca --- /dev/null +++ b/checkpoint/orbax/checkpoint/experimental/v1/_src/testing/compatibility/managed_checkpoints/v0_managed/has_metrics/no_custom_metadata/0/_CHECKPOINT_METADATA @@ -0,0 +1 @@ +{"item_handlers": {"baz": "orbax.checkpoint.experimental.v1._src.testing.handler_utils.DataclassCheckpointHandler", "metadata": "orbax.checkpoint._src.handlers.json_checkpoint_handler.JsonCheckpointHandler", "metrics": "orbax.checkpoint._src.handlers.json_checkpoint_handler.JsonCheckpointHandler", "state": "orbax.checkpoint._src.handlers.pytree_checkpoint_handler.PyTreeCheckpointHandler"}, "metrics": {}, "performance_metrics": {}, "init_timestamp_nsecs": 1778139891457166308, "commit_timestamp_nsecs": 1778139892772923073, "custom_metadata": {}} \ No newline at end of file diff --git a/checkpoint/orbax/checkpoint/experimental/v1/_src/testing/compatibility/managed_checkpoints/v0_managed/has_metrics/no_custom_metadata/0/baz/foo.txt b/checkpoint/orbax/checkpoint/experimental/v1/_src/testing/compatibility/managed_checkpoints/v0_managed/has_metrics/no_custom_metadata/0/baz/foo.txt new file mode 100644 index 000000000..c9b962c5a --- /dev/null +++ b/checkpoint/orbax/checkpoint/experimental/v1/_src/testing/compatibility/managed_checkpoints/v0_managed/has_metrics/no_custom_metadata/0/baz/foo.txt @@ -0,0 +1 @@ +{"int_val": 123, "str_val": "hi"} \ No newline at end of file diff --git a/checkpoint/orbax/checkpoint/experimental/v1/_src/testing/compatibility/managed_checkpoints/v0_managed/has_metrics/no_custom_metadata/0/metadata/metadata b/checkpoint/orbax/checkpoint/experimental/v1/_src/testing/compatibility/managed_checkpoints/v0_managed/has_metrics/no_custom_metadata/0/metadata/metadata new file mode 100644 index 000000000..7ac5d3905 --- /dev/null +++ b/checkpoint/orbax/checkpoint/experimental/v1/_src/testing/compatibility/managed_checkpoints/v0_managed/has_metrics/no_custom_metadata/0/metadata/metadata @@ -0,0 +1 @@ +{"metadata": "json_data"} \ No newline at end of file diff --git a/checkpoint/orbax/checkpoint/experimental/v1/_src/testing/compatibility/managed_checkpoints/v0_managed/has_metrics/no_custom_metadata/0/metrics/metrics b/checkpoint/orbax/checkpoint/experimental/v1/_src/testing/compatibility/managed_checkpoints/v0_managed/has_metrics/no_custom_metadata/0/metrics/metrics new file mode 100644 index 000000000..07c48241d --- /dev/null +++ b/checkpoint/orbax/checkpoint/experimental/v1/_src/testing/compatibility/managed_checkpoints/v0_managed/has_metrics/no_custom_metadata/0/metrics/metrics @@ -0,0 +1 @@ +{"loss": 0.5} \ No newline at end of file diff --git a/checkpoint/orbax/checkpoint/experimental/v1/_src/testing/compatibility/managed_checkpoints/v0_managed/has_metrics/no_custom_metadata/0/state/_METADATA b/checkpoint/orbax/checkpoint/experimental/v1/_src/testing/compatibility/managed_checkpoints/v0_managed/has_metrics/no_custom_metadata/0/state/_METADATA new file mode 100644 index 000000000..d18275eb9 --- /dev/null +++ b/checkpoint/orbax/checkpoint/experimental/v1/_src/testing/compatibility/managed_checkpoints/v0_managed/has_metrics/no_custom_metadata/0/state/_METADATA @@ -0,0 +1 @@ +{"tree_metadata": {"('a',)": {"key_metadata": [{"key": "a", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8]}}, "('b', 'c')": {"key_metadata": [{"key": "b", "key_type": 2}, {"key": "c", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3]}}}, "use_ocdbt": true, "use_zarr3": false, "store_array_data_equal_to_fill_value": true, "custom_metadata": null} \ No newline at end of file diff --git a/checkpoint/orbax/checkpoint/experimental/v1/_src/testing/compatibility/managed_checkpoints/v0_managed/has_metrics/no_custom_metadata/0/state/_sharding b/checkpoint/orbax/checkpoint/experimental/v1/_src/testing/compatibility/managed_checkpoints/v0_managed/has_metrics/no_custom_metadata/0/state/_sharding new file mode 100644 index 000000000..57439c751 --- /dev/null +++ b/checkpoint/orbax/checkpoint/experimental/v1/_src/testing/compatibility/managed_checkpoints/v0_managed/has_metrics/no_custom_metadata/0/state/_sharding @@ -0,0 +1 @@ +{"YQ==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cpu:0\"}","Yi5j":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cpu:0\"}"} \ No newline at end of file diff --git a/checkpoint/orbax/checkpoint/experimental/v1/_src/testing/compatibility/managed_checkpoints/v0_managed/has_metrics/no_custom_metadata/0/state/array_metadatas/process_0 b/checkpoint/orbax/checkpoint/experimental/v1/_src/testing/compatibility/managed_checkpoints/v0_managed/has_metrics/no_custom_metadata/0/state/array_metadatas/process_0 new file mode 100644 index 000000000..9e1980540 --- /dev/null +++ b/checkpoint/orbax/checkpoint/experimental/v1/_src/testing/compatibility/managed_checkpoints/v0_managed/has_metrics/no_custom_metadata/0/state/array_metadatas/process_0 @@ -0,0 +1 @@ +{"array_metadatas": [{"array_metadata": {"param_name": "a", "write_shape": [8], "chunk_shape": [8], "ext_metadata": null}}, {"array_metadata": {"param_name": "b.c", "write_shape": [3], "chunk_shape": [3], "ext_metadata": null}}]} \ No newline at end of file diff --git a/checkpoint/orbax/checkpoint/experimental/v1/_src/testing/compatibility/managed_checkpoints/v0_managed/has_metrics/no_custom_metadata/0/state/d/e8df828f773a05f06eee72e0a1ad0539 b/checkpoint/orbax/checkpoint/experimental/v1/_src/testing/compatibility/managed_checkpoints/v0_managed/has_metrics/no_custom_metadata/0/state/d/e8df828f773a05f06eee72e0a1ad0539 new file mode 100644 index 000000000..57b9216da Binary files /dev/null and b/checkpoint/orbax/checkpoint/experimental/v1/_src/testing/compatibility/managed_checkpoints/v0_managed/has_metrics/no_custom_metadata/0/state/d/e8df828f773a05f06eee72e0a1ad0539 differ diff --git a/checkpoint/orbax/checkpoint/experimental/v1/_src/testing/compatibility/managed_checkpoints/v0_managed/has_metrics/no_custom_metadata/0/state/descriptor/descriptor.pbtxt b/checkpoint/orbax/checkpoint/experimental/v1/_src/testing/compatibility/managed_checkpoints/v0_managed/has_metrics/no_custom_metadata/0/state/descriptor/descriptor.pbtxt new file mode 100644 index 000000000..33cb78056 --- /dev/null +++ b/checkpoint/orbax/checkpoint/experimental/v1/_src/testing/compatibility/managed_checkpoints/v0_managed/has_metrics/no_custom_metadata/0/state/descriptor/descriptor.pbtxt @@ -0,0 +1,17 @@ +uuid: "c29218c5-cc67-4e65-9384-3fbb6119ec28" +build_data { + timestamp { + seconds: 1778139795 + } + user: "angelmau" + hostname: "otbfg4.prod.google.com" + path: "/google/src/cloud/buildrabbit-username/buildrabbit-client/" + target: "//orbax/checkpoint/experimental/v1/_src/testing/compatibility:generate_v0_checkpoint_manager_checkpoints" + invocation_id: "a580b822-1e3a-4824-b47a-9bd884c066c7" + changelist: 911781978 + baseline_cl: 911781978 + workspace_id: "angelmau/345866" + citc_snapshot: 679 + client_status: BUILD_CLIENT_STATUS_UNSPECIFIED + verifiable: NON_VERIFIABLE +} diff --git a/checkpoint/orbax/checkpoint/experimental/v1/_src/testing/compatibility/managed_checkpoints/v0_managed/has_metrics/no_custom_metadata/0/state/descriptor/uuid-c29218c5-cc67-4e65-9384-3fbb6119ec28 b/checkpoint/orbax/checkpoint/experimental/v1/_src/testing/compatibility/managed_checkpoints/v0_managed/has_metrics/no_custom_metadata/0/state/descriptor/uuid-c29218c5-cc67-4e65-9384-3fbb6119ec28 new file mode 100644 index 000000000..fc8181fe8 --- /dev/null +++ b/checkpoint/orbax/checkpoint/experimental/v1/_src/testing/compatibility/managed_checkpoints/v0_managed/has_metrics/no_custom_metadata/0/state/descriptor/uuid-c29218c5-cc67-4e65-9384-3fbb6119ec28 @@ -0,0 +1 @@ +uuid: "c29218c5-cc67-4e65-9384-3fbb6119ec28" diff --git a/checkpoint/orbax/checkpoint/experimental/v1/_src/testing/compatibility/managed_checkpoints/v0_managed/has_metrics/no_custom_metadata/0/state/manifest.ocdbt b/checkpoint/orbax/checkpoint/experimental/v1/_src/testing/compatibility/managed_checkpoints/v0_managed/has_metrics/no_custom_metadata/0/state/manifest.ocdbt new file mode 100644 index 000000000..ca457e51a Binary files /dev/null and b/checkpoint/orbax/checkpoint/experimental/v1/_src/testing/compatibility/managed_checkpoints/v0_managed/has_metrics/no_custom_metadata/0/state/manifest.ocdbt differ diff --git a/checkpoint/orbax/checkpoint/experimental/v1/_src/testing/compatibility/managed_checkpoints/v0_managed/has_metrics/no_custom_metadata/0/state/ocdbt.process_0/d/018ab640ddaf1fdb3258b3291f8d342b b/checkpoint/orbax/checkpoint/experimental/v1/_src/testing/compatibility/managed_checkpoints/v0_managed/has_metrics/no_custom_metadata/0/state/ocdbt.process_0/d/018ab640ddaf1fdb3258b3291f8d342b new file mode 100644 index 000000000..8d6d8647a Binary files /dev/null and b/checkpoint/orbax/checkpoint/experimental/v1/_src/testing/compatibility/managed_checkpoints/v0_managed/has_metrics/no_custom_metadata/0/state/ocdbt.process_0/d/018ab640ddaf1fdb3258b3291f8d342b differ diff --git a/checkpoint/orbax/checkpoint/experimental/v1/_src/testing/compatibility/managed_checkpoints/v0_managed/has_metrics/no_custom_metadata/0/state/ocdbt.process_0/d/db6c229caac63f4d9254071e22ce2f72 b/checkpoint/orbax/checkpoint/experimental/v1/_src/testing/compatibility/managed_checkpoints/v0_managed/has_metrics/no_custom_metadata/0/state/ocdbt.process_0/d/db6c229caac63f4d9254071e22ce2f72 new file mode 100644 index 000000000..57b9216da Binary files /dev/null and b/checkpoint/orbax/checkpoint/experimental/v1/_src/testing/compatibility/managed_checkpoints/v0_managed/has_metrics/no_custom_metadata/0/state/ocdbt.process_0/d/db6c229caac63f4d9254071e22ce2f72 differ diff --git a/checkpoint/orbax/checkpoint/experimental/v1/_src/testing/compatibility/managed_checkpoints/v0_managed/has_metrics/no_custom_metadata/0/state/ocdbt.process_0/manifest.ocdbt b/checkpoint/orbax/checkpoint/experimental/v1/_src/testing/compatibility/managed_checkpoints/v0_managed/has_metrics/no_custom_metadata/0/state/ocdbt.process_0/manifest.ocdbt new file mode 100644 index 000000000..0e5cd23ed Binary files /dev/null and b/checkpoint/orbax/checkpoint/experimental/v1/_src/testing/compatibility/managed_checkpoints/v0_managed/has_metrics/no_custom_metadata/0/state/ocdbt.process_0/manifest.ocdbt differ diff --git a/checkpoint/orbax/checkpoint/experimental/v1/_src/testing/compatibility/managed_checkpoints/v0_managed/no_metrics/has_custom_metadata/0/_CHECKPOINT_METADATA b/checkpoint/orbax/checkpoint/experimental/v1/_src/testing/compatibility/managed_checkpoints/v0_managed/no_metrics/has_custom_metadata/0/_CHECKPOINT_METADATA new file mode 100644 index 000000000..e206bcc4a --- /dev/null +++ b/checkpoint/orbax/checkpoint/experimental/v1/_src/testing/compatibility/managed_checkpoints/v0_managed/no_metrics/has_custom_metadata/0/_CHECKPOINT_METADATA @@ -0,0 +1 @@ +{"item_handlers": {"baz": "orbax.checkpoint.experimental.v1._src.testing.handler_utils.DataclassCheckpointHandler", "metadata": "orbax.checkpoint._src.handlers.json_checkpoint_handler.JsonCheckpointHandler", "state": "orbax.checkpoint._src.handlers.pytree_checkpoint_handler.PyTreeCheckpointHandler"}, "metrics": {}, "performance_metrics": {}, "init_timestamp_nsecs": 1778139894804626172, "commit_timestamp_nsecs": 1778139895692318692, "custom_metadata": {"custom": "meta"}} \ No newline at end of file diff --git a/checkpoint/orbax/checkpoint/experimental/v1/_src/testing/compatibility/managed_checkpoints/v0_managed/no_metrics/has_custom_metadata/0/baz/foo.txt b/checkpoint/orbax/checkpoint/experimental/v1/_src/testing/compatibility/managed_checkpoints/v0_managed/no_metrics/has_custom_metadata/0/baz/foo.txt new file mode 100644 index 000000000..c9b962c5a --- /dev/null +++ b/checkpoint/orbax/checkpoint/experimental/v1/_src/testing/compatibility/managed_checkpoints/v0_managed/no_metrics/has_custom_metadata/0/baz/foo.txt @@ -0,0 +1 @@ +{"int_val": 123, "str_val": "hi"} \ No newline at end of file diff --git a/checkpoint/orbax/checkpoint/experimental/v1/_src/testing/compatibility/managed_checkpoints/v0_managed/no_metrics/has_custom_metadata/0/metadata/metadata b/checkpoint/orbax/checkpoint/experimental/v1/_src/testing/compatibility/managed_checkpoints/v0_managed/no_metrics/has_custom_metadata/0/metadata/metadata new file mode 100644 index 000000000..7ac5d3905 --- /dev/null +++ b/checkpoint/orbax/checkpoint/experimental/v1/_src/testing/compatibility/managed_checkpoints/v0_managed/no_metrics/has_custom_metadata/0/metadata/metadata @@ -0,0 +1 @@ +{"metadata": "json_data"} \ No newline at end of file diff --git a/checkpoint/orbax/checkpoint/experimental/v1/_src/testing/compatibility/managed_checkpoints/v0_managed/no_metrics/has_custom_metadata/0/state/_METADATA b/checkpoint/orbax/checkpoint/experimental/v1/_src/testing/compatibility/managed_checkpoints/v0_managed/no_metrics/has_custom_metadata/0/state/_METADATA new file mode 100644 index 000000000..d18275eb9 --- /dev/null +++ b/checkpoint/orbax/checkpoint/experimental/v1/_src/testing/compatibility/managed_checkpoints/v0_managed/no_metrics/has_custom_metadata/0/state/_METADATA @@ -0,0 +1 @@ +{"tree_metadata": {"('a',)": {"key_metadata": [{"key": "a", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8]}}, "('b', 'c')": {"key_metadata": [{"key": "b", "key_type": 2}, {"key": "c", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3]}}}, "use_ocdbt": true, "use_zarr3": false, "store_array_data_equal_to_fill_value": true, "custom_metadata": null} \ No newline at end of file diff --git a/checkpoint/orbax/checkpoint/experimental/v1/_src/testing/compatibility/managed_checkpoints/v0_managed/no_metrics/has_custom_metadata/0/state/_sharding b/checkpoint/orbax/checkpoint/experimental/v1/_src/testing/compatibility/managed_checkpoints/v0_managed/no_metrics/has_custom_metadata/0/state/_sharding new file mode 100644 index 000000000..57439c751 --- /dev/null +++ b/checkpoint/orbax/checkpoint/experimental/v1/_src/testing/compatibility/managed_checkpoints/v0_managed/no_metrics/has_custom_metadata/0/state/_sharding @@ -0,0 +1 @@ +{"YQ==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cpu:0\"}","Yi5j":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cpu:0\"}"} \ No newline at end of file diff --git a/checkpoint/orbax/checkpoint/experimental/v1/_src/testing/compatibility/managed_checkpoints/v0_managed/no_metrics/has_custom_metadata/0/state/array_metadatas/process_0 b/checkpoint/orbax/checkpoint/experimental/v1/_src/testing/compatibility/managed_checkpoints/v0_managed/no_metrics/has_custom_metadata/0/state/array_metadatas/process_0 new file mode 100644 index 000000000..9e1980540 --- /dev/null +++ b/checkpoint/orbax/checkpoint/experimental/v1/_src/testing/compatibility/managed_checkpoints/v0_managed/no_metrics/has_custom_metadata/0/state/array_metadatas/process_0 @@ -0,0 +1 @@ +{"array_metadatas": [{"array_metadata": {"param_name": "a", "write_shape": [8], "chunk_shape": [8], "ext_metadata": null}}, {"array_metadata": {"param_name": "b.c", "write_shape": [3], "chunk_shape": [3], "ext_metadata": null}}]} \ No newline at end of file diff --git a/checkpoint/orbax/checkpoint/experimental/v1/_src/testing/compatibility/managed_checkpoints/v0_managed/no_metrics/has_custom_metadata/0/state/d/e4ce8c428a2582cf4cdf06738cc3e428 b/checkpoint/orbax/checkpoint/experimental/v1/_src/testing/compatibility/managed_checkpoints/v0_managed/no_metrics/has_custom_metadata/0/state/d/e4ce8c428a2582cf4cdf06738cc3e428 new file mode 100644 index 000000000..57b9216da Binary files /dev/null and b/checkpoint/orbax/checkpoint/experimental/v1/_src/testing/compatibility/managed_checkpoints/v0_managed/no_metrics/has_custom_metadata/0/state/d/e4ce8c428a2582cf4cdf06738cc3e428 differ diff --git a/checkpoint/orbax/checkpoint/experimental/v1/_src/testing/compatibility/managed_checkpoints/v0_managed/no_metrics/has_custom_metadata/0/state/descriptor/descriptor.pbtxt b/checkpoint/orbax/checkpoint/experimental/v1/_src/testing/compatibility/managed_checkpoints/v0_managed/no_metrics/has_custom_metadata/0/state/descriptor/descriptor.pbtxt new file mode 100644 index 000000000..f673f3f13 --- /dev/null +++ b/checkpoint/orbax/checkpoint/experimental/v1/_src/testing/compatibility/managed_checkpoints/v0_managed/no_metrics/has_custom_metadata/0/state/descriptor/descriptor.pbtxt @@ -0,0 +1,17 @@ +uuid: "c1461d12-0db6-4b00-83c9-aa74486c53e6" +build_data { + timestamp { + seconds: 1778139795 + } + user: "angelmau" + hostname: "otbfg4.prod.google.com" + path: "/google/src/cloud/buildrabbit-username/buildrabbit-client/" + target: "//orbax/checkpoint/experimental/v1/_src/testing/compatibility:generate_v0_checkpoint_manager_checkpoints" + invocation_id: "a580b822-1e3a-4824-b47a-9bd884c066c7" + changelist: 911781978 + baseline_cl: 911781978 + workspace_id: "angelmau/345866" + citc_snapshot: 679 + client_status: BUILD_CLIENT_STATUS_UNSPECIFIED + verifiable: NON_VERIFIABLE +} diff --git a/checkpoint/orbax/checkpoint/experimental/v1/_src/testing/compatibility/managed_checkpoints/v0_managed/no_metrics/has_custom_metadata/0/state/descriptor/uuid-c1461d12-0db6-4b00-83c9-aa74486c53e6 b/checkpoint/orbax/checkpoint/experimental/v1/_src/testing/compatibility/managed_checkpoints/v0_managed/no_metrics/has_custom_metadata/0/state/descriptor/uuid-c1461d12-0db6-4b00-83c9-aa74486c53e6 new file mode 100644 index 000000000..caa5b9aea --- /dev/null +++ b/checkpoint/orbax/checkpoint/experimental/v1/_src/testing/compatibility/managed_checkpoints/v0_managed/no_metrics/has_custom_metadata/0/state/descriptor/uuid-c1461d12-0db6-4b00-83c9-aa74486c53e6 @@ -0,0 +1 @@ +uuid: "c1461d12-0db6-4b00-83c9-aa74486c53e6" diff --git a/checkpoint/orbax/checkpoint/experimental/v1/_src/testing/compatibility/managed_checkpoints/v0_managed/no_metrics/has_custom_metadata/0/state/manifest.ocdbt b/checkpoint/orbax/checkpoint/experimental/v1/_src/testing/compatibility/managed_checkpoints/v0_managed/no_metrics/has_custom_metadata/0/state/manifest.ocdbt new file mode 100644 index 000000000..3bf69308a Binary files /dev/null and b/checkpoint/orbax/checkpoint/experimental/v1/_src/testing/compatibility/managed_checkpoints/v0_managed/no_metrics/has_custom_metadata/0/state/manifest.ocdbt differ diff --git a/checkpoint/orbax/checkpoint/experimental/v1/_src/testing/compatibility/managed_checkpoints/v0_managed/no_metrics/has_custom_metadata/0/state/ocdbt.process_0/d/2eee718ac1419a5d32aa79f16581e1f5 b/checkpoint/orbax/checkpoint/experimental/v1/_src/testing/compatibility/managed_checkpoints/v0_managed/no_metrics/has_custom_metadata/0/state/ocdbt.process_0/d/2eee718ac1419a5d32aa79f16581e1f5 new file mode 100644 index 000000000..8d6d8647a Binary files /dev/null and b/checkpoint/orbax/checkpoint/experimental/v1/_src/testing/compatibility/managed_checkpoints/v0_managed/no_metrics/has_custom_metadata/0/state/ocdbt.process_0/d/2eee718ac1419a5d32aa79f16581e1f5 differ diff --git a/checkpoint/orbax/checkpoint/experimental/v1/_src/testing/compatibility/managed_checkpoints/v0_managed/no_metrics/has_custom_metadata/0/state/ocdbt.process_0/d/8e88ce9c0c365753a7c3ef9e0a99716f b/checkpoint/orbax/checkpoint/experimental/v1/_src/testing/compatibility/managed_checkpoints/v0_managed/no_metrics/has_custom_metadata/0/state/ocdbt.process_0/d/8e88ce9c0c365753a7c3ef9e0a99716f new file mode 100644 index 000000000..a8adc2748 Binary files /dev/null and b/checkpoint/orbax/checkpoint/experimental/v1/_src/testing/compatibility/managed_checkpoints/v0_managed/no_metrics/has_custom_metadata/0/state/ocdbt.process_0/d/8e88ce9c0c365753a7c3ef9e0a99716f differ diff --git a/checkpoint/orbax/checkpoint/experimental/v1/_src/testing/compatibility/managed_checkpoints/v0_managed/no_metrics/has_custom_metadata/0/state/ocdbt.process_0/d/a97a2190f0a7f8a20f1e3ba20ae0761f b/checkpoint/orbax/checkpoint/experimental/v1/_src/testing/compatibility/managed_checkpoints/v0_managed/no_metrics/has_custom_metadata/0/state/ocdbt.process_0/d/a97a2190f0a7f8a20f1e3ba20ae0761f new file mode 100644 index 000000000..57b9216da Binary files /dev/null and b/checkpoint/orbax/checkpoint/experimental/v1/_src/testing/compatibility/managed_checkpoints/v0_managed/no_metrics/has_custom_metadata/0/state/ocdbt.process_0/d/a97a2190f0a7f8a20f1e3ba20ae0761f differ diff --git a/checkpoint/orbax/checkpoint/experimental/v1/_src/testing/compatibility/managed_checkpoints/v0_managed/no_metrics/has_custom_metadata/0/state/ocdbt.process_0/d/db5bdead14946c3b48f3259e8b2c26fa b/checkpoint/orbax/checkpoint/experimental/v1/_src/testing/compatibility/managed_checkpoints/v0_managed/no_metrics/has_custom_metadata/0/state/ocdbt.process_0/d/db5bdead14946c3b48f3259e8b2c26fa new file mode 100644 index 000000000..34ec758d3 Binary files /dev/null and b/checkpoint/orbax/checkpoint/experimental/v1/_src/testing/compatibility/managed_checkpoints/v0_managed/no_metrics/has_custom_metadata/0/state/ocdbt.process_0/d/db5bdead14946c3b48f3259e8b2c26fa differ diff --git a/checkpoint/orbax/checkpoint/experimental/v1/_src/testing/compatibility/managed_checkpoints/v0_managed/no_metrics/has_custom_metadata/0/state/ocdbt.process_0/manifest.ocdbt b/checkpoint/orbax/checkpoint/experimental/v1/_src/testing/compatibility/managed_checkpoints/v0_managed/no_metrics/has_custom_metadata/0/state/ocdbt.process_0/manifest.ocdbt new file mode 100644 index 000000000..be7ceac5a Binary files /dev/null and b/checkpoint/orbax/checkpoint/experimental/v1/_src/testing/compatibility/managed_checkpoints/v0_managed/no_metrics/has_custom_metadata/0/state/ocdbt.process_0/manifest.ocdbt differ diff --git a/checkpoint/orbax/checkpoint/experimental/v1/_src/testing/compatibility/managed_checkpoints/v0_managed/no_metrics/has_custom_metadata/metadata/_ROOT_METADATA b/checkpoint/orbax/checkpoint/experimental/v1/_src/testing/compatibility/managed_checkpoints/v0_managed/no_metrics/has_custom_metadata/metadata/_ROOT_METADATA new file mode 100644 index 000000000..0b4608849 --- /dev/null +++ b/checkpoint/orbax/checkpoint/experimental/v1/_src/testing/compatibility/managed_checkpoints/v0_managed/no_metrics/has_custom_metadata/metadata/_ROOT_METADATA @@ -0,0 +1 @@ +{"internal_metadata": {}, "custom_metadata": {"foo": "bar"}} \ No newline at end of file diff --git a/checkpoint/orbax/checkpoint/experimental/v1/_src/testing/compatibility/managed_checkpoints/v0_managed/no_metrics/no_custom_metadata/0/_CHECKPOINT_METADATA b/checkpoint/orbax/checkpoint/experimental/v1/_src/testing/compatibility/managed_checkpoints/v0_managed/no_metrics/no_custom_metadata/0/_CHECKPOINT_METADATA new file mode 100644 index 000000000..ad58acfde --- /dev/null +++ b/checkpoint/orbax/checkpoint/experimental/v1/_src/testing/compatibility/managed_checkpoints/v0_managed/no_metrics/no_custom_metadata/0/_CHECKPOINT_METADATA @@ -0,0 +1 @@ +{"item_handlers": {"baz": "orbax.checkpoint.experimental.v1._src.testing.handler_utils.DataclassCheckpointHandler", "metadata": "orbax.checkpoint._src.handlers.json_checkpoint_handler.JsonCheckpointHandler", "state": "orbax.checkpoint._src.handlers.pytree_checkpoint_handler.PyTreeCheckpointHandler"}, "metrics": {}, "performance_metrics": {}, "init_timestamp_nsecs": 1778139896170460708, "commit_timestamp_nsecs": 1778139897332055157, "custom_metadata": {}} \ No newline at end of file diff --git a/checkpoint/orbax/checkpoint/experimental/v1/_src/testing/compatibility/managed_checkpoints/v0_managed/no_metrics/no_custom_metadata/0/baz/foo.txt b/checkpoint/orbax/checkpoint/experimental/v1/_src/testing/compatibility/managed_checkpoints/v0_managed/no_metrics/no_custom_metadata/0/baz/foo.txt new file mode 100644 index 000000000..c9b962c5a --- /dev/null +++ b/checkpoint/orbax/checkpoint/experimental/v1/_src/testing/compatibility/managed_checkpoints/v0_managed/no_metrics/no_custom_metadata/0/baz/foo.txt @@ -0,0 +1 @@ +{"int_val": 123, "str_val": "hi"} \ No newline at end of file diff --git a/checkpoint/orbax/checkpoint/experimental/v1/_src/testing/compatibility/managed_checkpoints/v0_managed/no_metrics/no_custom_metadata/0/metadata/metadata b/checkpoint/orbax/checkpoint/experimental/v1/_src/testing/compatibility/managed_checkpoints/v0_managed/no_metrics/no_custom_metadata/0/metadata/metadata new file mode 100644 index 000000000..7ac5d3905 --- /dev/null +++ b/checkpoint/orbax/checkpoint/experimental/v1/_src/testing/compatibility/managed_checkpoints/v0_managed/no_metrics/no_custom_metadata/0/metadata/metadata @@ -0,0 +1 @@ +{"metadata": "json_data"} \ No newline at end of file diff --git a/checkpoint/orbax/checkpoint/experimental/v1/_src/testing/compatibility/managed_checkpoints/v0_managed/no_metrics/no_custom_metadata/0/state/_METADATA b/checkpoint/orbax/checkpoint/experimental/v1/_src/testing/compatibility/managed_checkpoints/v0_managed/no_metrics/no_custom_metadata/0/state/_METADATA new file mode 100644 index 000000000..d18275eb9 --- /dev/null +++ b/checkpoint/orbax/checkpoint/experimental/v1/_src/testing/compatibility/managed_checkpoints/v0_managed/no_metrics/no_custom_metadata/0/state/_METADATA @@ -0,0 +1 @@ +{"tree_metadata": {"('a',)": {"key_metadata": [{"key": "a", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8]}}, "('b', 'c')": {"key_metadata": [{"key": "b", "key_type": 2}, {"key": "c", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3]}}}, "use_ocdbt": true, "use_zarr3": false, "store_array_data_equal_to_fill_value": true, "custom_metadata": null} \ No newline at end of file diff --git a/checkpoint/orbax/checkpoint/experimental/v1/_src/testing/compatibility/managed_checkpoints/v0_managed/no_metrics/no_custom_metadata/0/state/_sharding b/checkpoint/orbax/checkpoint/experimental/v1/_src/testing/compatibility/managed_checkpoints/v0_managed/no_metrics/no_custom_metadata/0/state/_sharding new file mode 100644 index 000000000..57439c751 --- /dev/null +++ b/checkpoint/orbax/checkpoint/experimental/v1/_src/testing/compatibility/managed_checkpoints/v0_managed/no_metrics/no_custom_metadata/0/state/_sharding @@ -0,0 +1 @@ +{"YQ==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cpu:0\"}","Yi5j":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cpu:0\"}"} \ No newline at end of file diff --git a/checkpoint/orbax/checkpoint/experimental/v1/_src/testing/compatibility/managed_checkpoints/v0_managed/no_metrics/no_custom_metadata/0/state/array_metadatas/process_0 b/checkpoint/orbax/checkpoint/experimental/v1/_src/testing/compatibility/managed_checkpoints/v0_managed/no_metrics/no_custom_metadata/0/state/array_metadatas/process_0 new file mode 100644 index 000000000..9e1980540 --- /dev/null +++ b/checkpoint/orbax/checkpoint/experimental/v1/_src/testing/compatibility/managed_checkpoints/v0_managed/no_metrics/no_custom_metadata/0/state/array_metadatas/process_0 @@ -0,0 +1 @@ +{"array_metadatas": [{"array_metadata": {"param_name": "a", "write_shape": [8], "chunk_shape": [8], "ext_metadata": null}}, {"array_metadata": {"param_name": "b.c", "write_shape": [3], "chunk_shape": [3], "ext_metadata": null}}]} \ No newline at end of file diff --git a/checkpoint/orbax/checkpoint/experimental/v1/_src/testing/compatibility/managed_checkpoints/v0_managed/no_metrics/no_custom_metadata/0/state/d/b0f28f119de0006d0f365a3463cca553 b/checkpoint/orbax/checkpoint/experimental/v1/_src/testing/compatibility/managed_checkpoints/v0_managed/no_metrics/no_custom_metadata/0/state/d/b0f28f119de0006d0f365a3463cca553 new file mode 100644 index 000000000..57b9216da Binary files /dev/null and b/checkpoint/orbax/checkpoint/experimental/v1/_src/testing/compatibility/managed_checkpoints/v0_managed/no_metrics/no_custom_metadata/0/state/d/b0f28f119de0006d0f365a3463cca553 differ diff --git a/checkpoint/orbax/checkpoint/experimental/v1/_src/testing/compatibility/managed_checkpoints/v0_managed/no_metrics/no_custom_metadata/0/state/descriptor/descriptor.pbtxt b/checkpoint/orbax/checkpoint/experimental/v1/_src/testing/compatibility/managed_checkpoints/v0_managed/no_metrics/no_custom_metadata/0/state/descriptor/descriptor.pbtxt new file mode 100644 index 000000000..2351a401e --- /dev/null +++ b/checkpoint/orbax/checkpoint/experimental/v1/_src/testing/compatibility/managed_checkpoints/v0_managed/no_metrics/no_custom_metadata/0/state/descriptor/descriptor.pbtxt @@ -0,0 +1,17 @@ +uuid: "dd7e3917-b5bd-43ae-9aa0-b1af1b48694f" +build_data { + timestamp { + seconds: 1778139795 + } + user: "angelmau" + hostname: "otbfg4.prod.google.com" + path: "/google/src/cloud/buildrabbit-username/buildrabbit-client/" + target: "//orbax/checkpoint/experimental/v1/_src/testing/compatibility:generate_v0_checkpoint_manager_checkpoints" + invocation_id: "a580b822-1e3a-4824-b47a-9bd884c066c7" + changelist: 911781978 + baseline_cl: 911781978 + workspace_id: "angelmau/345866" + citc_snapshot: 679 + client_status: BUILD_CLIENT_STATUS_UNSPECIFIED + verifiable: NON_VERIFIABLE +} diff --git a/checkpoint/orbax/checkpoint/experimental/v1/_src/testing/compatibility/managed_checkpoints/v0_managed/no_metrics/no_custom_metadata/0/state/descriptor/uuid-dd7e3917-b5bd-43ae-9aa0-b1af1b48694f b/checkpoint/orbax/checkpoint/experimental/v1/_src/testing/compatibility/managed_checkpoints/v0_managed/no_metrics/no_custom_metadata/0/state/descriptor/uuid-dd7e3917-b5bd-43ae-9aa0-b1af1b48694f new file mode 100644 index 000000000..ca7c23737 --- /dev/null +++ b/checkpoint/orbax/checkpoint/experimental/v1/_src/testing/compatibility/managed_checkpoints/v0_managed/no_metrics/no_custom_metadata/0/state/descriptor/uuid-dd7e3917-b5bd-43ae-9aa0-b1af1b48694f @@ -0,0 +1 @@ +uuid: "dd7e3917-b5bd-43ae-9aa0-b1af1b48694f" diff --git a/checkpoint/orbax/checkpoint/experimental/v1/_src/testing/compatibility/managed_checkpoints/v0_managed/no_metrics/no_custom_metadata/0/state/manifest.ocdbt b/checkpoint/orbax/checkpoint/experimental/v1/_src/testing/compatibility/managed_checkpoints/v0_managed/no_metrics/no_custom_metadata/0/state/manifest.ocdbt new file mode 100644 index 000000000..41d5c57df Binary files /dev/null and b/checkpoint/orbax/checkpoint/experimental/v1/_src/testing/compatibility/managed_checkpoints/v0_managed/no_metrics/no_custom_metadata/0/state/manifest.ocdbt differ diff --git a/checkpoint/orbax/checkpoint/experimental/v1/_src/testing/compatibility/managed_checkpoints/v0_managed/no_metrics/no_custom_metadata/0/state/ocdbt.process_0/d/16fec959e6da1e856f1c06ebe52298f8 b/checkpoint/orbax/checkpoint/experimental/v1/_src/testing/compatibility/managed_checkpoints/v0_managed/no_metrics/no_custom_metadata/0/state/ocdbt.process_0/d/16fec959e6da1e856f1c06ebe52298f8 new file mode 100644 index 000000000..63a24bbda Binary files /dev/null and b/checkpoint/orbax/checkpoint/experimental/v1/_src/testing/compatibility/managed_checkpoints/v0_managed/no_metrics/no_custom_metadata/0/state/ocdbt.process_0/d/16fec959e6da1e856f1c06ebe52298f8 differ diff --git a/checkpoint/orbax/checkpoint/experimental/v1/_src/testing/compatibility/managed_checkpoints/v0_managed/no_metrics/no_custom_metadata/0/state/ocdbt.process_0/d/3aa68a447e9b8b913343fa39c56b6618 b/checkpoint/orbax/checkpoint/experimental/v1/_src/testing/compatibility/managed_checkpoints/v0_managed/no_metrics/no_custom_metadata/0/state/ocdbt.process_0/d/3aa68a447e9b8b913343fa39c56b6618 new file mode 100644 index 000000000..8d6d8647a Binary files /dev/null and b/checkpoint/orbax/checkpoint/experimental/v1/_src/testing/compatibility/managed_checkpoints/v0_managed/no_metrics/no_custom_metadata/0/state/ocdbt.process_0/d/3aa68a447e9b8b913343fa39c56b6618 differ diff --git a/checkpoint/orbax/checkpoint/experimental/v1/_src/testing/compatibility/managed_checkpoints/v0_managed/no_metrics/no_custom_metadata/0/state/ocdbt.process_0/d/c900507fa78eb4adfa5f14ac0d0de3ec b/checkpoint/orbax/checkpoint/experimental/v1/_src/testing/compatibility/managed_checkpoints/v0_managed/no_metrics/no_custom_metadata/0/state/ocdbt.process_0/d/c900507fa78eb4adfa5f14ac0d0de3ec new file mode 100644 index 000000000..317f53f94 Binary files /dev/null and b/checkpoint/orbax/checkpoint/experimental/v1/_src/testing/compatibility/managed_checkpoints/v0_managed/no_metrics/no_custom_metadata/0/state/ocdbt.process_0/d/c900507fa78eb4adfa5f14ac0d0de3ec differ diff --git a/checkpoint/orbax/checkpoint/experimental/v1/_src/testing/compatibility/managed_checkpoints/v0_managed/no_metrics/no_custom_metadata/0/state/ocdbt.process_0/d/f74f53bddc8092376c216e00bd1abeb6 b/checkpoint/orbax/checkpoint/experimental/v1/_src/testing/compatibility/managed_checkpoints/v0_managed/no_metrics/no_custom_metadata/0/state/ocdbt.process_0/d/f74f53bddc8092376c216e00bd1abeb6 new file mode 100644 index 000000000..57b9216da Binary files /dev/null and b/checkpoint/orbax/checkpoint/experimental/v1/_src/testing/compatibility/managed_checkpoints/v0_managed/no_metrics/no_custom_metadata/0/state/ocdbt.process_0/d/f74f53bddc8092376c216e00bd1abeb6 differ diff --git a/checkpoint/orbax/checkpoint/experimental/v1/_src/testing/compatibility/managed_checkpoints/v0_managed/no_metrics/no_custom_metadata/0/state/ocdbt.process_0/manifest.ocdbt b/checkpoint/orbax/checkpoint/experimental/v1/_src/testing/compatibility/managed_checkpoints/v0_managed/no_metrics/no_custom_metadata/0/state/ocdbt.process_0/manifest.ocdbt new file mode 100644 index 000000000..5e4acf20b Binary files /dev/null and b/checkpoint/orbax/checkpoint/experimental/v1/_src/testing/compatibility/managed_checkpoints/v0_managed/no_metrics/no_custom_metadata/0/state/ocdbt.process_0/manifest.ocdbt differ diff --git a/checkpoint/orbax/checkpoint/experimental/v1/_src/testing/compatibility/managed_checkpoints/v1_managed/has_metrics/has_custom_metadata/0/_CHECKPOINT_METADATA b/checkpoint/orbax/checkpoint/experimental/v1/_src/testing/compatibility/managed_checkpoints/v1_managed/has_metrics/has_custom_metadata/0/_CHECKPOINT_METADATA new file mode 100644 index 000000000..df53a5fe3 --- /dev/null +++ b/checkpoint/orbax/checkpoint/experimental/v1/_src/testing/compatibility/managed_checkpoints/v1_managed/has_metrics/has_custom_metadata/0/_CHECKPOINT_METADATA @@ -0,0 +1 @@ +{"item_handlers": {"baz": "orbax.checkpoint.experimental.v1._src.testing.handler_utils.BazHandler", "descriptor": "orbax.checkpoint.experimental.v1._src.handlers.descriptor_handler.DescriptorHandler", "metadata": "orbax.checkpoint.experimental.v1._src.handlers.json_handler.JsonHandler", "metrics": "orbax.checkpoint.experimental.v1._src.handlers.json_handler.MetricsHandler", "state": "orbax.checkpoint.experimental.v1._src.handlers.pytree_handler.PyTreeHandler"}, "metrics": {}, "performance_metrics": {}, "init_timestamp_nsecs": 1778139945096880714, "commit_timestamp_nsecs": 1778139947624156729, "custom_metadata": {"custom": "meta"}} \ No newline at end of file diff --git a/checkpoint/orbax/checkpoint/experimental/v1/_src/testing/compatibility/managed_checkpoints/v1_managed/has_metrics/has_custom_metadata/0/baz/foo.txt b/checkpoint/orbax/checkpoint/experimental/v1/_src/testing/compatibility/managed_checkpoints/v1_managed/has_metrics/has_custom_metadata/0/baz/foo.txt new file mode 100644 index 000000000..c9b962c5a --- /dev/null +++ b/checkpoint/orbax/checkpoint/experimental/v1/_src/testing/compatibility/managed_checkpoints/v1_managed/has_metrics/has_custom_metadata/0/baz/foo.txt @@ -0,0 +1 @@ +{"int_val": 123, "str_val": "hi"} \ No newline at end of file diff --git a/checkpoint/orbax/checkpoint/experimental/v1/_src/testing/compatibility/managed_checkpoints/v1_managed/has_metrics/has_custom_metadata/0/descriptor/descriptor.pbtxt b/checkpoint/orbax/checkpoint/experimental/v1/_src/testing/compatibility/managed_checkpoints/v1_managed/has_metrics/has_custom_metadata/0/descriptor/descriptor.pbtxt new file mode 100644 index 000000000..32ac2dc21 --- /dev/null +++ b/checkpoint/orbax/checkpoint/experimental/v1/_src/testing/compatibility/managed_checkpoints/v1_managed/has_metrics/has_custom_metadata/0/descriptor/descriptor.pbtxt @@ -0,0 +1,17 @@ +uuid: "8e9fc4b4-0a03-4f12-b342-e5364a856703" +build_data { + timestamp { + seconds: 1778139916 + } + user: "angelmau" + hostname: "otbfg4.prod.google.com" + path: "/google/src/cloud/buildrabbit-username/buildrabbit-client/" + target: "//orbax/checkpoint/experimental/v1/_src/testing/compatibility:generate_v1_checkpointer_checkpoints" + invocation_id: "cdd06442-0c11-4760-8e7a-b6c7d35b454e" + changelist: 911781978 + baseline_cl: 911781978 + workspace_id: "angelmau/345866" + citc_snapshot: 701 + client_status: BUILD_CLIENT_STATUS_UNSPECIFIED + verifiable: NON_VERIFIABLE +} diff --git a/checkpoint/orbax/checkpoint/experimental/v1/_src/testing/compatibility/managed_checkpoints/v1_managed/has_metrics/has_custom_metadata/0/descriptor/uuid-8e9fc4b4-0a03-4f12-b342-e5364a856703 b/checkpoint/orbax/checkpoint/experimental/v1/_src/testing/compatibility/managed_checkpoints/v1_managed/has_metrics/has_custom_metadata/0/descriptor/uuid-8e9fc4b4-0a03-4f12-b342-e5364a856703 new file mode 100644 index 000000000..a700dc484 --- /dev/null +++ b/checkpoint/orbax/checkpoint/experimental/v1/_src/testing/compatibility/managed_checkpoints/v1_managed/has_metrics/has_custom_metadata/0/descriptor/uuid-8e9fc4b4-0a03-4f12-b342-e5364a856703 @@ -0,0 +1 @@ +uuid: "8e9fc4b4-0a03-4f12-b342-e5364a856703" diff --git a/checkpoint/orbax/checkpoint/experimental/v1/_src/testing/compatibility/managed_checkpoints/v1_managed/has_metrics/has_custom_metadata/0/metadata/data.json b/checkpoint/orbax/checkpoint/experimental/v1/_src/testing/compatibility/managed_checkpoints/v1_managed/has_metrics/has_custom_metadata/0/metadata/data.json new file mode 100644 index 000000000..7ac5d3905 --- /dev/null +++ b/checkpoint/orbax/checkpoint/experimental/v1/_src/testing/compatibility/managed_checkpoints/v1_managed/has_metrics/has_custom_metadata/0/metadata/data.json @@ -0,0 +1 @@ +{"metadata": "json_data"} \ No newline at end of file diff --git a/checkpoint/orbax/checkpoint/experimental/v1/_src/testing/compatibility/managed_checkpoints/v1_managed/has_metrics/has_custom_metadata/0/metrics/metrics b/checkpoint/orbax/checkpoint/experimental/v1/_src/testing/compatibility/managed_checkpoints/v1_managed/has_metrics/has_custom_metadata/0/metrics/metrics new file mode 100644 index 000000000..07c48241d --- /dev/null +++ b/checkpoint/orbax/checkpoint/experimental/v1/_src/testing/compatibility/managed_checkpoints/v1_managed/has_metrics/has_custom_metadata/0/metrics/metrics @@ -0,0 +1 @@ +{"loss": 0.5} \ No newline at end of file diff --git a/checkpoint/orbax/checkpoint/experimental/v1/_src/testing/compatibility/managed_checkpoints/v1_managed/has_metrics/has_custom_metadata/0/state/_METADATA b/checkpoint/orbax/checkpoint/experimental/v1/_src/testing/compatibility/managed_checkpoints/v1_managed/has_metrics/has_custom_metadata/0/state/_METADATA new file mode 100644 index 000000000..58dab14f3 --- /dev/null +++ b/checkpoint/orbax/checkpoint/experimental/v1/_src/testing/compatibility/managed_checkpoints/v1_managed/has_metrics/has_custom_metadata/0/state/_METADATA @@ -0,0 +1 @@ +{"tree_metadata": {"('a',)": {"key_metadata": [{"key": "a", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8]}}, "('b', 'c')": {"key_metadata": [{"key": "b", "key_type": 2}, {"key": "c", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3]}}}, "use_ocdbt": true, "use_zarr3": true, "store_array_data_equal_to_fill_value": true, "custom_metadata": null} \ No newline at end of file diff --git a/checkpoint/orbax/checkpoint/experimental/v1/_src/testing/compatibility/managed_checkpoints/v1_managed/has_metrics/has_custom_metadata/0/state/_sharding b/checkpoint/orbax/checkpoint/experimental/v1/_src/testing/compatibility/managed_checkpoints/v1_managed/has_metrics/has_custom_metadata/0/state/_sharding new file mode 100644 index 000000000..57439c751 --- /dev/null +++ b/checkpoint/orbax/checkpoint/experimental/v1/_src/testing/compatibility/managed_checkpoints/v1_managed/has_metrics/has_custom_metadata/0/state/_sharding @@ -0,0 +1 @@ +{"YQ==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cpu:0\"}","Yi5j":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cpu:0\"}"} \ No newline at end of file diff --git a/checkpoint/orbax/checkpoint/experimental/v1/_src/testing/compatibility/managed_checkpoints/v1_managed/has_metrics/has_custom_metadata/0/state/array_metadatas/process_0 b/checkpoint/orbax/checkpoint/experimental/v1/_src/testing/compatibility/managed_checkpoints/v1_managed/has_metrics/has_custom_metadata/0/state/array_metadatas/process_0 new file mode 100644 index 000000000..9e1980540 --- /dev/null +++ b/checkpoint/orbax/checkpoint/experimental/v1/_src/testing/compatibility/managed_checkpoints/v1_managed/has_metrics/has_custom_metadata/0/state/array_metadatas/process_0 @@ -0,0 +1 @@ +{"array_metadatas": [{"array_metadata": {"param_name": "a", "write_shape": [8], "chunk_shape": [8], "ext_metadata": null}}, {"array_metadata": {"param_name": "b.c", "write_shape": [3], "chunk_shape": [3], "ext_metadata": null}}]} \ No newline at end of file diff --git a/checkpoint/orbax/checkpoint/experimental/v1/_src/testing/compatibility/managed_checkpoints/v1_managed/has_metrics/has_custom_metadata/0/state/d/480e6bd7514fd2ae2f11da6a51437c28 b/checkpoint/orbax/checkpoint/experimental/v1/_src/testing/compatibility/managed_checkpoints/v1_managed/has_metrics/has_custom_metadata/0/state/d/480e6bd7514fd2ae2f11da6a51437c28 new file mode 100644 index 000000000..dc3b09d1b Binary files /dev/null and b/checkpoint/orbax/checkpoint/experimental/v1/_src/testing/compatibility/managed_checkpoints/v1_managed/has_metrics/has_custom_metadata/0/state/d/480e6bd7514fd2ae2f11da6a51437c28 differ diff --git a/checkpoint/orbax/checkpoint/experimental/v1/_src/testing/compatibility/managed_checkpoints/v1_managed/has_metrics/has_custom_metadata/0/state/manifest.ocdbt b/checkpoint/orbax/checkpoint/experimental/v1/_src/testing/compatibility/managed_checkpoints/v1_managed/has_metrics/has_custom_metadata/0/state/manifest.ocdbt new file mode 100644 index 000000000..d2eba8887 Binary files /dev/null and b/checkpoint/orbax/checkpoint/experimental/v1/_src/testing/compatibility/managed_checkpoints/v1_managed/has_metrics/has_custom_metadata/0/state/manifest.ocdbt differ diff --git a/checkpoint/orbax/checkpoint/experimental/v1/_src/testing/compatibility/managed_checkpoints/v1_managed/has_metrics/has_custom_metadata/0/state/ocdbt.process_0/d/8e18a0958239c8ecf1895b06b1318a01 b/checkpoint/orbax/checkpoint/experimental/v1/_src/testing/compatibility/managed_checkpoints/v1_managed/has_metrics/has_custom_metadata/0/state/ocdbt.process_0/d/8e18a0958239c8ecf1895b06b1318a01 new file mode 100644 index 000000000..0f5426c04 Binary files /dev/null and b/checkpoint/orbax/checkpoint/experimental/v1/_src/testing/compatibility/managed_checkpoints/v1_managed/has_metrics/has_custom_metadata/0/state/ocdbt.process_0/d/8e18a0958239c8ecf1895b06b1318a01 differ diff --git a/checkpoint/orbax/checkpoint/experimental/v1/_src/testing/compatibility/managed_checkpoints/v1_managed/has_metrics/has_custom_metadata/0/state/ocdbt.process_0/d/a2427bf1070aa986a1826fb23a78654f b/checkpoint/orbax/checkpoint/experimental/v1/_src/testing/compatibility/managed_checkpoints/v1_managed/has_metrics/has_custom_metadata/0/state/ocdbt.process_0/d/a2427bf1070aa986a1826fb23a78654f new file mode 100644 index 000000000..dc3b09d1b Binary files /dev/null and b/checkpoint/orbax/checkpoint/experimental/v1/_src/testing/compatibility/managed_checkpoints/v1_managed/has_metrics/has_custom_metadata/0/state/ocdbt.process_0/d/a2427bf1070aa986a1826fb23a78654f differ diff --git a/checkpoint/orbax/checkpoint/experimental/v1/_src/testing/compatibility/managed_checkpoints/v1_managed/has_metrics/has_custom_metadata/0/state/ocdbt.process_0/d/f86fde568f5080ee1c2b2adfe13562a8 b/checkpoint/orbax/checkpoint/experimental/v1/_src/testing/compatibility/managed_checkpoints/v1_managed/has_metrics/has_custom_metadata/0/state/ocdbt.process_0/d/f86fde568f5080ee1c2b2adfe13562a8 new file mode 100644 index 000000000..d610a0fd3 Binary files /dev/null and b/checkpoint/orbax/checkpoint/experimental/v1/_src/testing/compatibility/managed_checkpoints/v1_managed/has_metrics/has_custom_metadata/0/state/ocdbt.process_0/d/f86fde568f5080ee1c2b2adfe13562a8 differ diff --git a/checkpoint/orbax/checkpoint/experimental/v1/_src/testing/compatibility/managed_checkpoints/v1_managed/has_metrics/has_custom_metadata/0/state/ocdbt.process_0/manifest.ocdbt b/checkpoint/orbax/checkpoint/experimental/v1/_src/testing/compatibility/managed_checkpoints/v1_managed/has_metrics/has_custom_metadata/0/state/ocdbt.process_0/manifest.ocdbt new file mode 100644 index 000000000..3770a746a Binary files /dev/null and b/checkpoint/orbax/checkpoint/experimental/v1/_src/testing/compatibility/managed_checkpoints/v1_managed/has_metrics/has_custom_metadata/0/state/ocdbt.process_0/manifest.ocdbt differ diff --git a/checkpoint/orbax/checkpoint/experimental/v1/_src/testing/compatibility/managed_checkpoints/v1_managed/has_metrics/has_custom_metadata/metadata/_ROOT_METADATA b/checkpoint/orbax/checkpoint/experimental/v1/_src/testing/compatibility/managed_checkpoints/v1_managed/has_metrics/has_custom_metadata/metadata/_ROOT_METADATA new file mode 100644 index 000000000..0b4608849 --- /dev/null +++ b/checkpoint/orbax/checkpoint/experimental/v1/_src/testing/compatibility/managed_checkpoints/v1_managed/has_metrics/has_custom_metadata/metadata/_ROOT_METADATA @@ -0,0 +1 @@ +{"internal_metadata": {}, "custom_metadata": {"foo": "bar"}} \ No newline at end of file diff --git a/checkpoint/orbax/checkpoint/experimental/v1/_src/testing/compatibility/managed_checkpoints/v1_managed/has_metrics/no_custom_metadata/0/_CHECKPOINT_METADATA b/checkpoint/orbax/checkpoint/experimental/v1/_src/testing/compatibility/managed_checkpoints/v1_managed/has_metrics/no_custom_metadata/0/_CHECKPOINT_METADATA new file mode 100644 index 000000000..d2c57443e --- /dev/null +++ b/checkpoint/orbax/checkpoint/experimental/v1/_src/testing/compatibility/managed_checkpoints/v1_managed/has_metrics/no_custom_metadata/0/_CHECKPOINT_METADATA @@ -0,0 +1 @@ +{"item_handlers": {"baz": "orbax.checkpoint.experimental.v1._src.testing.handler_utils.BazHandler", "descriptor": "orbax.checkpoint.experimental.v1._src.handlers.descriptor_handler.DescriptorHandler", "metadata": "orbax.checkpoint.experimental.v1._src.handlers.json_handler.JsonHandler", "metrics": "orbax.checkpoint.experimental.v1._src.handlers.json_handler.MetricsHandler", "state": "orbax.checkpoint.experimental.v1._src.handlers.pytree_handler.PyTreeHandler"}, "metrics": {}, "performance_metrics": {}, "init_timestamp_nsecs": 1778139948562443946, "commit_timestamp_nsecs": 1778139950826891462, "custom_metadata": {}} \ No newline at end of file diff --git a/checkpoint/orbax/checkpoint/experimental/v1/_src/testing/compatibility/managed_checkpoints/v1_managed/has_metrics/no_custom_metadata/0/baz/foo.txt b/checkpoint/orbax/checkpoint/experimental/v1/_src/testing/compatibility/managed_checkpoints/v1_managed/has_metrics/no_custom_metadata/0/baz/foo.txt new file mode 100644 index 000000000..c9b962c5a --- /dev/null +++ b/checkpoint/orbax/checkpoint/experimental/v1/_src/testing/compatibility/managed_checkpoints/v1_managed/has_metrics/no_custom_metadata/0/baz/foo.txt @@ -0,0 +1 @@ +{"int_val": 123, "str_val": "hi"} \ No newline at end of file diff --git a/checkpoint/orbax/checkpoint/experimental/v1/_src/testing/compatibility/managed_checkpoints/v1_managed/has_metrics/no_custom_metadata/0/descriptor/descriptor.pbtxt b/checkpoint/orbax/checkpoint/experimental/v1/_src/testing/compatibility/managed_checkpoints/v1_managed/has_metrics/no_custom_metadata/0/descriptor/descriptor.pbtxt new file mode 100644 index 000000000..d12d69ab3 --- /dev/null +++ b/checkpoint/orbax/checkpoint/experimental/v1/_src/testing/compatibility/managed_checkpoints/v1_managed/has_metrics/no_custom_metadata/0/descriptor/descriptor.pbtxt @@ -0,0 +1,17 @@ +uuid: "508df3ba-9bee-4777-bcfc-e8054d780976" +build_data { + timestamp { + seconds: 1778139916 + } + user: "angelmau" + hostname: "otbfg4.prod.google.com" + path: "/google/src/cloud/buildrabbit-username/buildrabbit-client/" + target: "//orbax/checkpoint/experimental/v1/_src/testing/compatibility:generate_v1_checkpointer_checkpoints" + invocation_id: "cdd06442-0c11-4760-8e7a-b6c7d35b454e" + changelist: 911781978 + baseline_cl: 911781978 + workspace_id: "angelmau/345866" + citc_snapshot: 701 + client_status: BUILD_CLIENT_STATUS_UNSPECIFIED + verifiable: NON_VERIFIABLE +} diff --git a/checkpoint/orbax/checkpoint/experimental/v1/_src/testing/compatibility/managed_checkpoints/v1_managed/has_metrics/no_custom_metadata/0/descriptor/uuid-508df3ba-9bee-4777-bcfc-e8054d780976 b/checkpoint/orbax/checkpoint/experimental/v1/_src/testing/compatibility/managed_checkpoints/v1_managed/has_metrics/no_custom_metadata/0/descriptor/uuid-508df3ba-9bee-4777-bcfc-e8054d780976 new file mode 100644 index 000000000..8df3e20de --- /dev/null +++ b/checkpoint/orbax/checkpoint/experimental/v1/_src/testing/compatibility/managed_checkpoints/v1_managed/has_metrics/no_custom_metadata/0/descriptor/uuid-508df3ba-9bee-4777-bcfc-e8054d780976 @@ -0,0 +1 @@ +uuid: "508df3ba-9bee-4777-bcfc-e8054d780976" diff --git a/checkpoint/orbax/checkpoint/experimental/v1/_src/testing/compatibility/managed_checkpoints/v1_managed/has_metrics/no_custom_metadata/0/metadata/data.json b/checkpoint/orbax/checkpoint/experimental/v1/_src/testing/compatibility/managed_checkpoints/v1_managed/has_metrics/no_custom_metadata/0/metadata/data.json new file mode 100644 index 000000000..7ac5d3905 --- /dev/null +++ b/checkpoint/orbax/checkpoint/experimental/v1/_src/testing/compatibility/managed_checkpoints/v1_managed/has_metrics/no_custom_metadata/0/metadata/data.json @@ -0,0 +1 @@ +{"metadata": "json_data"} \ No newline at end of file diff --git a/checkpoint/orbax/checkpoint/experimental/v1/_src/testing/compatibility/managed_checkpoints/v1_managed/has_metrics/no_custom_metadata/0/metrics/metrics b/checkpoint/orbax/checkpoint/experimental/v1/_src/testing/compatibility/managed_checkpoints/v1_managed/has_metrics/no_custom_metadata/0/metrics/metrics new file mode 100644 index 000000000..07c48241d --- /dev/null +++ b/checkpoint/orbax/checkpoint/experimental/v1/_src/testing/compatibility/managed_checkpoints/v1_managed/has_metrics/no_custom_metadata/0/metrics/metrics @@ -0,0 +1 @@ +{"loss": 0.5} \ No newline at end of file diff --git a/checkpoint/orbax/checkpoint/experimental/v1/_src/testing/compatibility/managed_checkpoints/v1_managed/has_metrics/no_custom_metadata/0/state/_METADATA b/checkpoint/orbax/checkpoint/experimental/v1/_src/testing/compatibility/managed_checkpoints/v1_managed/has_metrics/no_custom_metadata/0/state/_METADATA new file mode 100644 index 000000000..58dab14f3 --- /dev/null +++ b/checkpoint/orbax/checkpoint/experimental/v1/_src/testing/compatibility/managed_checkpoints/v1_managed/has_metrics/no_custom_metadata/0/state/_METADATA @@ -0,0 +1 @@ +{"tree_metadata": {"('a',)": {"key_metadata": [{"key": "a", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8]}}, "('b', 'c')": {"key_metadata": [{"key": "b", "key_type": 2}, {"key": "c", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3]}}}, "use_ocdbt": true, "use_zarr3": true, "store_array_data_equal_to_fill_value": true, "custom_metadata": null} \ No newline at end of file diff --git a/checkpoint/orbax/checkpoint/experimental/v1/_src/testing/compatibility/managed_checkpoints/v1_managed/has_metrics/no_custom_metadata/0/state/_sharding b/checkpoint/orbax/checkpoint/experimental/v1/_src/testing/compatibility/managed_checkpoints/v1_managed/has_metrics/no_custom_metadata/0/state/_sharding new file mode 100644 index 000000000..57439c751 --- /dev/null +++ b/checkpoint/orbax/checkpoint/experimental/v1/_src/testing/compatibility/managed_checkpoints/v1_managed/has_metrics/no_custom_metadata/0/state/_sharding @@ -0,0 +1 @@ +{"YQ==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cpu:0\"}","Yi5j":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cpu:0\"}"} \ No newline at end of file diff --git a/checkpoint/orbax/checkpoint/experimental/v1/_src/testing/compatibility/managed_checkpoints/v1_managed/has_metrics/no_custom_metadata/0/state/array_metadatas/process_0 b/checkpoint/orbax/checkpoint/experimental/v1/_src/testing/compatibility/managed_checkpoints/v1_managed/has_metrics/no_custom_metadata/0/state/array_metadatas/process_0 new file mode 100644 index 000000000..9e1980540 --- /dev/null +++ b/checkpoint/orbax/checkpoint/experimental/v1/_src/testing/compatibility/managed_checkpoints/v1_managed/has_metrics/no_custom_metadata/0/state/array_metadatas/process_0 @@ -0,0 +1 @@ +{"array_metadatas": [{"array_metadata": {"param_name": "a", "write_shape": [8], "chunk_shape": [8], "ext_metadata": null}}, {"array_metadata": {"param_name": "b.c", "write_shape": [3], "chunk_shape": [3], "ext_metadata": null}}]} \ No newline at end of file diff --git a/checkpoint/orbax/checkpoint/experimental/v1/_src/testing/compatibility/managed_checkpoints/v1_managed/has_metrics/no_custom_metadata/0/state/d/8d9a25bc839918b4fcf92cf52a7124c9 b/checkpoint/orbax/checkpoint/experimental/v1/_src/testing/compatibility/managed_checkpoints/v1_managed/has_metrics/no_custom_metadata/0/state/d/8d9a25bc839918b4fcf92cf52a7124c9 new file mode 100644 index 000000000..dc3b09d1b Binary files /dev/null and b/checkpoint/orbax/checkpoint/experimental/v1/_src/testing/compatibility/managed_checkpoints/v1_managed/has_metrics/no_custom_metadata/0/state/d/8d9a25bc839918b4fcf92cf52a7124c9 differ diff --git a/checkpoint/orbax/checkpoint/experimental/v1/_src/testing/compatibility/managed_checkpoints/v1_managed/has_metrics/no_custom_metadata/0/state/manifest.ocdbt b/checkpoint/orbax/checkpoint/experimental/v1/_src/testing/compatibility/managed_checkpoints/v1_managed/has_metrics/no_custom_metadata/0/state/manifest.ocdbt new file mode 100644 index 000000000..34a0321af Binary files /dev/null and b/checkpoint/orbax/checkpoint/experimental/v1/_src/testing/compatibility/managed_checkpoints/v1_managed/has_metrics/no_custom_metadata/0/state/manifest.ocdbt differ diff --git a/checkpoint/orbax/checkpoint/experimental/v1/_src/testing/compatibility/managed_checkpoints/v1_managed/has_metrics/no_custom_metadata/0/state/ocdbt.process_0/d/1494480e45d4ee3eaedc0bd7737d5ab4 b/checkpoint/orbax/checkpoint/experimental/v1/_src/testing/compatibility/managed_checkpoints/v1_managed/has_metrics/no_custom_metadata/0/state/ocdbt.process_0/d/1494480e45d4ee3eaedc0bd7737d5ab4 new file mode 100644 index 000000000..d610a0fd3 Binary files /dev/null and b/checkpoint/orbax/checkpoint/experimental/v1/_src/testing/compatibility/managed_checkpoints/v1_managed/has_metrics/no_custom_metadata/0/state/ocdbt.process_0/d/1494480e45d4ee3eaedc0bd7737d5ab4 differ diff --git a/checkpoint/orbax/checkpoint/experimental/v1/_src/testing/compatibility/managed_checkpoints/v1_managed/has_metrics/no_custom_metadata/0/state/ocdbt.process_0/d/cb04bc578eee9b0493e04285e07c2369 b/checkpoint/orbax/checkpoint/experimental/v1/_src/testing/compatibility/managed_checkpoints/v1_managed/has_metrics/no_custom_metadata/0/state/ocdbt.process_0/d/cb04bc578eee9b0493e04285e07c2369 new file mode 100644 index 000000000..0f5426c04 Binary files /dev/null and b/checkpoint/orbax/checkpoint/experimental/v1/_src/testing/compatibility/managed_checkpoints/v1_managed/has_metrics/no_custom_metadata/0/state/ocdbt.process_0/d/cb04bc578eee9b0493e04285e07c2369 differ diff --git a/checkpoint/orbax/checkpoint/experimental/v1/_src/testing/compatibility/managed_checkpoints/v1_managed/has_metrics/no_custom_metadata/0/state/ocdbt.process_0/d/d0b0a0fcf359d15408bf78f3b2b62945 b/checkpoint/orbax/checkpoint/experimental/v1/_src/testing/compatibility/managed_checkpoints/v1_managed/has_metrics/no_custom_metadata/0/state/ocdbt.process_0/d/d0b0a0fcf359d15408bf78f3b2b62945 new file mode 100644 index 000000000..dc3b09d1b Binary files /dev/null and b/checkpoint/orbax/checkpoint/experimental/v1/_src/testing/compatibility/managed_checkpoints/v1_managed/has_metrics/no_custom_metadata/0/state/ocdbt.process_0/d/d0b0a0fcf359d15408bf78f3b2b62945 differ diff --git a/checkpoint/orbax/checkpoint/experimental/v1/_src/testing/compatibility/managed_checkpoints/v1_managed/has_metrics/no_custom_metadata/0/state/ocdbt.process_0/manifest.ocdbt b/checkpoint/orbax/checkpoint/experimental/v1/_src/testing/compatibility/managed_checkpoints/v1_managed/has_metrics/no_custom_metadata/0/state/ocdbt.process_0/manifest.ocdbt new file mode 100644 index 000000000..2e49f05dd Binary files /dev/null and b/checkpoint/orbax/checkpoint/experimental/v1/_src/testing/compatibility/managed_checkpoints/v1_managed/has_metrics/no_custom_metadata/0/state/ocdbt.process_0/manifest.ocdbt differ diff --git a/checkpoint/orbax/checkpoint/experimental/v1/_src/testing/compatibility/managed_checkpoints/v1_managed/no_metrics/has_custom_metadata/0/_CHECKPOINT_METADATA b/checkpoint/orbax/checkpoint/experimental/v1/_src/testing/compatibility/managed_checkpoints/v1_managed/no_metrics/has_custom_metadata/0/_CHECKPOINT_METADATA new file mode 100644 index 000000000..e503bf718 --- /dev/null +++ b/checkpoint/orbax/checkpoint/experimental/v1/_src/testing/compatibility/managed_checkpoints/v1_managed/no_metrics/has_custom_metadata/0/_CHECKPOINT_METADATA @@ -0,0 +1 @@ +{"item_handlers": {"baz": "orbax.checkpoint.experimental.v1._src.testing.handler_utils.BazHandler", "descriptor": "orbax.checkpoint.experimental.v1._src.handlers.descriptor_handler.DescriptorHandler", "metadata": "orbax.checkpoint.experimental.v1._src.handlers.json_handler.JsonHandler", "state": "orbax.checkpoint.experimental.v1._src.handlers.pytree_handler.PyTreeHandler"}, "metrics": {}, "performance_metrics": {}, "init_timestamp_nsecs": 1778139952850671788, "commit_timestamp_nsecs": 1778139954529877780, "custom_metadata": {"custom": "meta"}} \ No newline at end of file diff --git a/checkpoint/orbax/checkpoint/experimental/v1/_src/testing/compatibility/managed_checkpoints/v1_managed/no_metrics/has_custom_metadata/0/baz/foo.txt b/checkpoint/orbax/checkpoint/experimental/v1/_src/testing/compatibility/managed_checkpoints/v1_managed/no_metrics/has_custom_metadata/0/baz/foo.txt new file mode 100644 index 000000000..c9b962c5a --- /dev/null +++ b/checkpoint/orbax/checkpoint/experimental/v1/_src/testing/compatibility/managed_checkpoints/v1_managed/no_metrics/has_custom_metadata/0/baz/foo.txt @@ -0,0 +1 @@ +{"int_val": 123, "str_val": "hi"} \ No newline at end of file diff --git a/checkpoint/orbax/checkpoint/experimental/v1/_src/testing/compatibility/managed_checkpoints/v1_managed/no_metrics/has_custom_metadata/0/descriptor/descriptor.pbtxt b/checkpoint/orbax/checkpoint/experimental/v1/_src/testing/compatibility/managed_checkpoints/v1_managed/no_metrics/has_custom_metadata/0/descriptor/descriptor.pbtxt new file mode 100644 index 000000000..872644f9c --- /dev/null +++ b/checkpoint/orbax/checkpoint/experimental/v1/_src/testing/compatibility/managed_checkpoints/v1_managed/no_metrics/has_custom_metadata/0/descriptor/descriptor.pbtxt @@ -0,0 +1,17 @@ +uuid: "c3e829de-2ae2-41d2-9a5b-2120c3777b43" +build_data { + timestamp { + seconds: 1778139916 + } + user: "angelmau" + hostname: "otbfg4.prod.google.com" + path: "/google/src/cloud/buildrabbit-username/buildrabbit-client/" + target: "//orbax/checkpoint/experimental/v1/_src/testing/compatibility:generate_v1_checkpointer_checkpoints" + invocation_id: "cdd06442-0c11-4760-8e7a-b6c7d35b454e" + changelist: 911781978 + baseline_cl: 911781978 + workspace_id: "angelmau/345866" + citc_snapshot: 701 + client_status: BUILD_CLIENT_STATUS_UNSPECIFIED + verifiable: NON_VERIFIABLE +} diff --git a/checkpoint/orbax/checkpoint/experimental/v1/_src/testing/compatibility/managed_checkpoints/v1_managed/no_metrics/has_custom_metadata/0/descriptor/uuid-c3e829de-2ae2-41d2-9a5b-2120c3777b43 b/checkpoint/orbax/checkpoint/experimental/v1/_src/testing/compatibility/managed_checkpoints/v1_managed/no_metrics/has_custom_metadata/0/descriptor/uuid-c3e829de-2ae2-41d2-9a5b-2120c3777b43 new file mode 100644 index 000000000..3a2c4299c --- /dev/null +++ b/checkpoint/orbax/checkpoint/experimental/v1/_src/testing/compatibility/managed_checkpoints/v1_managed/no_metrics/has_custom_metadata/0/descriptor/uuid-c3e829de-2ae2-41d2-9a5b-2120c3777b43 @@ -0,0 +1 @@ +uuid: "c3e829de-2ae2-41d2-9a5b-2120c3777b43" diff --git a/checkpoint/orbax/checkpoint/experimental/v1/_src/testing/compatibility/managed_checkpoints/v1_managed/no_metrics/has_custom_metadata/0/metadata/data.json b/checkpoint/orbax/checkpoint/experimental/v1/_src/testing/compatibility/managed_checkpoints/v1_managed/no_metrics/has_custom_metadata/0/metadata/data.json new file mode 100644 index 000000000..7ac5d3905 --- /dev/null +++ b/checkpoint/orbax/checkpoint/experimental/v1/_src/testing/compatibility/managed_checkpoints/v1_managed/no_metrics/has_custom_metadata/0/metadata/data.json @@ -0,0 +1 @@ +{"metadata": "json_data"} \ No newline at end of file diff --git a/checkpoint/orbax/checkpoint/experimental/v1/_src/testing/compatibility/managed_checkpoints/v1_managed/no_metrics/has_custom_metadata/0/state/_METADATA b/checkpoint/orbax/checkpoint/experimental/v1/_src/testing/compatibility/managed_checkpoints/v1_managed/no_metrics/has_custom_metadata/0/state/_METADATA new file mode 100644 index 000000000..58dab14f3 --- /dev/null +++ b/checkpoint/orbax/checkpoint/experimental/v1/_src/testing/compatibility/managed_checkpoints/v1_managed/no_metrics/has_custom_metadata/0/state/_METADATA @@ -0,0 +1 @@ +{"tree_metadata": {"('a',)": {"key_metadata": [{"key": "a", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8]}}, "('b', 'c')": {"key_metadata": [{"key": "b", "key_type": 2}, {"key": "c", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3]}}}, "use_ocdbt": true, "use_zarr3": true, "store_array_data_equal_to_fill_value": true, "custom_metadata": null} \ No newline at end of file diff --git a/checkpoint/orbax/checkpoint/experimental/v1/_src/testing/compatibility/managed_checkpoints/v1_managed/no_metrics/has_custom_metadata/0/state/_sharding b/checkpoint/orbax/checkpoint/experimental/v1/_src/testing/compatibility/managed_checkpoints/v1_managed/no_metrics/has_custom_metadata/0/state/_sharding new file mode 100644 index 000000000..57439c751 --- /dev/null +++ b/checkpoint/orbax/checkpoint/experimental/v1/_src/testing/compatibility/managed_checkpoints/v1_managed/no_metrics/has_custom_metadata/0/state/_sharding @@ -0,0 +1 @@ +{"YQ==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cpu:0\"}","Yi5j":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cpu:0\"}"} \ No newline at end of file diff --git a/checkpoint/orbax/checkpoint/experimental/v1/_src/testing/compatibility/managed_checkpoints/v1_managed/no_metrics/has_custom_metadata/0/state/array_metadatas/process_0 b/checkpoint/orbax/checkpoint/experimental/v1/_src/testing/compatibility/managed_checkpoints/v1_managed/no_metrics/has_custom_metadata/0/state/array_metadatas/process_0 new file mode 100644 index 000000000..9e1980540 --- /dev/null +++ b/checkpoint/orbax/checkpoint/experimental/v1/_src/testing/compatibility/managed_checkpoints/v1_managed/no_metrics/has_custom_metadata/0/state/array_metadatas/process_0 @@ -0,0 +1 @@ +{"array_metadatas": [{"array_metadata": {"param_name": "a", "write_shape": [8], "chunk_shape": [8], "ext_metadata": null}}, {"array_metadata": {"param_name": "b.c", "write_shape": [3], "chunk_shape": [3], "ext_metadata": null}}]} \ No newline at end of file diff --git a/checkpoint/orbax/checkpoint/experimental/v1/_src/testing/compatibility/managed_checkpoints/v1_managed/no_metrics/has_custom_metadata/0/state/d/ac931c24315ca0ed0940335b3d254e78 b/checkpoint/orbax/checkpoint/experimental/v1/_src/testing/compatibility/managed_checkpoints/v1_managed/no_metrics/has_custom_metadata/0/state/d/ac931c24315ca0ed0940335b3d254e78 new file mode 100644 index 000000000..dc3b09d1b Binary files /dev/null and b/checkpoint/orbax/checkpoint/experimental/v1/_src/testing/compatibility/managed_checkpoints/v1_managed/no_metrics/has_custom_metadata/0/state/d/ac931c24315ca0ed0940335b3d254e78 differ diff --git a/checkpoint/orbax/checkpoint/experimental/v1/_src/testing/compatibility/managed_checkpoints/v1_managed/no_metrics/has_custom_metadata/0/state/manifest.ocdbt b/checkpoint/orbax/checkpoint/experimental/v1/_src/testing/compatibility/managed_checkpoints/v1_managed/no_metrics/has_custom_metadata/0/state/manifest.ocdbt new file mode 100644 index 000000000..0b6ad009b Binary files /dev/null and b/checkpoint/orbax/checkpoint/experimental/v1/_src/testing/compatibility/managed_checkpoints/v1_managed/no_metrics/has_custom_metadata/0/state/manifest.ocdbt differ diff --git a/checkpoint/orbax/checkpoint/experimental/v1/_src/testing/compatibility/managed_checkpoints/v1_managed/no_metrics/has_custom_metadata/0/state/ocdbt.process_0/d/1bdceb1a55e3b70569527962f465bcb3 b/checkpoint/orbax/checkpoint/experimental/v1/_src/testing/compatibility/managed_checkpoints/v1_managed/no_metrics/has_custom_metadata/0/state/ocdbt.process_0/d/1bdceb1a55e3b70569527962f465bcb3 new file mode 100644 index 000000000..d610a0fd3 Binary files /dev/null and b/checkpoint/orbax/checkpoint/experimental/v1/_src/testing/compatibility/managed_checkpoints/v1_managed/no_metrics/has_custom_metadata/0/state/ocdbt.process_0/d/1bdceb1a55e3b70569527962f465bcb3 differ diff --git a/checkpoint/orbax/checkpoint/experimental/v1/_src/testing/compatibility/managed_checkpoints/v1_managed/no_metrics/has_custom_metadata/0/state/ocdbt.process_0/d/5139eb6383ba81d772163c584015116f b/checkpoint/orbax/checkpoint/experimental/v1/_src/testing/compatibility/managed_checkpoints/v1_managed/no_metrics/has_custom_metadata/0/state/ocdbt.process_0/d/5139eb6383ba81d772163c584015116f new file mode 100644 index 000000000..dc3b09d1b Binary files /dev/null and b/checkpoint/orbax/checkpoint/experimental/v1/_src/testing/compatibility/managed_checkpoints/v1_managed/no_metrics/has_custom_metadata/0/state/ocdbt.process_0/d/5139eb6383ba81d772163c584015116f differ diff --git a/checkpoint/orbax/checkpoint/experimental/v1/_src/testing/compatibility/managed_checkpoints/v1_managed/no_metrics/has_custom_metadata/0/state/ocdbt.process_0/manifest.ocdbt b/checkpoint/orbax/checkpoint/experimental/v1/_src/testing/compatibility/managed_checkpoints/v1_managed/no_metrics/has_custom_metadata/0/state/ocdbt.process_0/manifest.ocdbt new file mode 100644 index 000000000..6dd2ee643 Binary files /dev/null and b/checkpoint/orbax/checkpoint/experimental/v1/_src/testing/compatibility/managed_checkpoints/v1_managed/no_metrics/has_custom_metadata/0/state/ocdbt.process_0/manifest.ocdbt differ diff --git a/checkpoint/orbax/checkpoint/experimental/v1/_src/testing/compatibility/managed_checkpoints/v1_managed/no_metrics/has_custom_metadata/metadata/_ROOT_METADATA b/checkpoint/orbax/checkpoint/experimental/v1/_src/testing/compatibility/managed_checkpoints/v1_managed/no_metrics/has_custom_metadata/metadata/_ROOT_METADATA new file mode 100644 index 000000000..0b4608849 --- /dev/null +++ b/checkpoint/orbax/checkpoint/experimental/v1/_src/testing/compatibility/managed_checkpoints/v1_managed/no_metrics/has_custom_metadata/metadata/_ROOT_METADATA @@ -0,0 +1 @@ +{"internal_metadata": {}, "custom_metadata": {"foo": "bar"}} \ No newline at end of file diff --git a/checkpoint/orbax/checkpoint/experimental/v1/_src/testing/compatibility/managed_checkpoints/v1_managed/no_metrics/no_custom_metadata/0/_CHECKPOINT_METADATA b/checkpoint/orbax/checkpoint/experimental/v1/_src/testing/compatibility/managed_checkpoints/v1_managed/no_metrics/no_custom_metadata/0/_CHECKPOINT_METADATA new file mode 100644 index 000000000..d6e8a320e --- /dev/null +++ b/checkpoint/orbax/checkpoint/experimental/v1/_src/testing/compatibility/managed_checkpoints/v1_managed/no_metrics/no_custom_metadata/0/_CHECKPOINT_METADATA @@ -0,0 +1 @@ +{"item_handlers": {"baz": "orbax.checkpoint.experimental.v1._src.testing.handler_utils.BazHandler", "descriptor": "orbax.checkpoint.experimental.v1._src.handlers.descriptor_handler.DescriptorHandler", "metadata": "orbax.checkpoint.experimental.v1._src.handlers.json_handler.JsonHandler", "state": "orbax.checkpoint.experimental.v1._src.handlers.pytree_handler.PyTreeHandler"}, "metrics": {}, "performance_metrics": {}, "init_timestamp_nsecs": 1778139955025172542, "commit_timestamp_nsecs": 1778139957541959436, "custom_metadata": {}} \ No newline at end of file diff --git a/checkpoint/orbax/checkpoint/experimental/v1/_src/testing/compatibility/managed_checkpoints/v1_managed/no_metrics/no_custom_metadata/0/baz/foo.txt b/checkpoint/orbax/checkpoint/experimental/v1/_src/testing/compatibility/managed_checkpoints/v1_managed/no_metrics/no_custom_metadata/0/baz/foo.txt new file mode 100644 index 000000000..c9b962c5a --- /dev/null +++ b/checkpoint/orbax/checkpoint/experimental/v1/_src/testing/compatibility/managed_checkpoints/v1_managed/no_metrics/no_custom_metadata/0/baz/foo.txt @@ -0,0 +1 @@ +{"int_val": 123, "str_val": "hi"} \ No newline at end of file diff --git a/checkpoint/orbax/checkpoint/experimental/v1/_src/testing/compatibility/managed_checkpoints/v1_managed/no_metrics/no_custom_metadata/0/descriptor/descriptor.pbtxt b/checkpoint/orbax/checkpoint/experimental/v1/_src/testing/compatibility/managed_checkpoints/v1_managed/no_metrics/no_custom_metadata/0/descriptor/descriptor.pbtxt new file mode 100644 index 000000000..0228bf260 --- /dev/null +++ b/checkpoint/orbax/checkpoint/experimental/v1/_src/testing/compatibility/managed_checkpoints/v1_managed/no_metrics/no_custom_metadata/0/descriptor/descriptor.pbtxt @@ -0,0 +1,17 @@ +uuid: "3d7ed91d-8849-4c13-804f-7c278ae22f4e" +build_data { + timestamp { + seconds: 1778139916 + } + user: "angelmau" + hostname: "otbfg4.prod.google.com" + path: "/google/src/cloud/buildrabbit-username/buildrabbit-client/" + target: "//orbax/checkpoint/experimental/v1/_src/testing/compatibility:generate_v1_checkpointer_checkpoints" + invocation_id: "cdd06442-0c11-4760-8e7a-b6c7d35b454e" + changelist: 911781978 + baseline_cl: 911781978 + workspace_id: "angelmau/345866" + citc_snapshot: 701 + client_status: BUILD_CLIENT_STATUS_UNSPECIFIED + verifiable: NON_VERIFIABLE +} diff --git a/checkpoint/orbax/checkpoint/experimental/v1/_src/testing/compatibility/managed_checkpoints/v1_managed/no_metrics/no_custom_metadata/0/descriptor/uuid-3d7ed91d-8849-4c13-804f-7c278ae22f4e b/checkpoint/orbax/checkpoint/experimental/v1/_src/testing/compatibility/managed_checkpoints/v1_managed/no_metrics/no_custom_metadata/0/descriptor/uuid-3d7ed91d-8849-4c13-804f-7c278ae22f4e new file mode 100644 index 000000000..473ab1b81 --- /dev/null +++ b/checkpoint/orbax/checkpoint/experimental/v1/_src/testing/compatibility/managed_checkpoints/v1_managed/no_metrics/no_custom_metadata/0/descriptor/uuid-3d7ed91d-8849-4c13-804f-7c278ae22f4e @@ -0,0 +1 @@ +uuid: "3d7ed91d-8849-4c13-804f-7c278ae22f4e" diff --git a/checkpoint/orbax/checkpoint/experimental/v1/_src/testing/compatibility/managed_checkpoints/v1_managed/no_metrics/no_custom_metadata/0/metadata/data.json b/checkpoint/orbax/checkpoint/experimental/v1/_src/testing/compatibility/managed_checkpoints/v1_managed/no_metrics/no_custom_metadata/0/metadata/data.json new file mode 100644 index 000000000..7ac5d3905 --- /dev/null +++ b/checkpoint/orbax/checkpoint/experimental/v1/_src/testing/compatibility/managed_checkpoints/v1_managed/no_metrics/no_custom_metadata/0/metadata/data.json @@ -0,0 +1 @@ +{"metadata": "json_data"} \ No newline at end of file diff --git a/checkpoint/orbax/checkpoint/experimental/v1/_src/testing/compatibility/managed_checkpoints/v1_managed/no_metrics/no_custom_metadata/0/state/_METADATA b/checkpoint/orbax/checkpoint/experimental/v1/_src/testing/compatibility/managed_checkpoints/v1_managed/no_metrics/no_custom_metadata/0/state/_METADATA new file mode 100644 index 000000000..58dab14f3 --- /dev/null +++ b/checkpoint/orbax/checkpoint/experimental/v1/_src/testing/compatibility/managed_checkpoints/v1_managed/no_metrics/no_custom_metadata/0/state/_METADATA @@ -0,0 +1 @@ +{"tree_metadata": {"('a',)": {"key_metadata": [{"key": "a", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8]}}, "('b', 'c')": {"key_metadata": [{"key": "b", "key_type": 2}, {"key": "c", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3]}}}, "use_ocdbt": true, "use_zarr3": true, "store_array_data_equal_to_fill_value": true, "custom_metadata": null} \ No newline at end of file diff --git a/checkpoint/orbax/checkpoint/experimental/v1/_src/testing/compatibility/managed_checkpoints/v1_managed/no_metrics/no_custom_metadata/0/state/_sharding b/checkpoint/orbax/checkpoint/experimental/v1/_src/testing/compatibility/managed_checkpoints/v1_managed/no_metrics/no_custom_metadata/0/state/_sharding new file mode 100644 index 000000000..57439c751 --- /dev/null +++ b/checkpoint/orbax/checkpoint/experimental/v1/_src/testing/compatibility/managed_checkpoints/v1_managed/no_metrics/no_custom_metadata/0/state/_sharding @@ -0,0 +1 @@ +{"YQ==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cpu:0\"}","Yi5j":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cpu:0\"}"} \ No newline at end of file diff --git a/checkpoint/orbax/checkpoint/experimental/v1/_src/testing/compatibility/managed_checkpoints/v1_managed/no_metrics/no_custom_metadata/0/state/array_metadatas/process_0 b/checkpoint/orbax/checkpoint/experimental/v1/_src/testing/compatibility/managed_checkpoints/v1_managed/no_metrics/no_custom_metadata/0/state/array_metadatas/process_0 new file mode 100644 index 000000000..9e1980540 --- /dev/null +++ b/checkpoint/orbax/checkpoint/experimental/v1/_src/testing/compatibility/managed_checkpoints/v1_managed/no_metrics/no_custom_metadata/0/state/array_metadatas/process_0 @@ -0,0 +1 @@ +{"array_metadatas": [{"array_metadata": {"param_name": "a", "write_shape": [8], "chunk_shape": [8], "ext_metadata": null}}, {"array_metadata": {"param_name": "b.c", "write_shape": [3], "chunk_shape": [3], "ext_metadata": null}}]} \ No newline at end of file diff --git a/checkpoint/orbax/checkpoint/experimental/v1/_src/testing/compatibility/managed_checkpoints/v1_managed/no_metrics/no_custom_metadata/0/state/d/d759cea92eb7825e071cdaf236421308 b/checkpoint/orbax/checkpoint/experimental/v1/_src/testing/compatibility/managed_checkpoints/v1_managed/no_metrics/no_custom_metadata/0/state/d/d759cea92eb7825e071cdaf236421308 new file mode 100644 index 000000000..dc3b09d1b Binary files /dev/null and b/checkpoint/orbax/checkpoint/experimental/v1/_src/testing/compatibility/managed_checkpoints/v1_managed/no_metrics/no_custom_metadata/0/state/d/d759cea92eb7825e071cdaf236421308 differ diff --git a/checkpoint/orbax/checkpoint/experimental/v1/_src/testing/compatibility/managed_checkpoints/v1_managed/no_metrics/no_custom_metadata/0/state/manifest.ocdbt b/checkpoint/orbax/checkpoint/experimental/v1/_src/testing/compatibility/managed_checkpoints/v1_managed/no_metrics/no_custom_metadata/0/state/manifest.ocdbt new file mode 100644 index 000000000..910d501ea Binary files /dev/null and b/checkpoint/orbax/checkpoint/experimental/v1/_src/testing/compatibility/managed_checkpoints/v1_managed/no_metrics/no_custom_metadata/0/state/manifest.ocdbt differ diff --git a/checkpoint/orbax/checkpoint/experimental/v1/_src/testing/compatibility/managed_checkpoints/v1_managed/no_metrics/no_custom_metadata/0/state/ocdbt.process_0/d/b47b7183ae4c1865684f50c448253b5a b/checkpoint/orbax/checkpoint/experimental/v1/_src/testing/compatibility/managed_checkpoints/v1_managed/no_metrics/no_custom_metadata/0/state/ocdbt.process_0/d/b47b7183ae4c1865684f50c448253b5a new file mode 100644 index 000000000..0f5426c04 Binary files /dev/null and b/checkpoint/orbax/checkpoint/experimental/v1/_src/testing/compatibility/managed_checkpoints/v1_managed/no_metrics/no_custom_metadata/0/state/ocdbt.process_0/d/b47b7183ae4c1865684f50c448253b5a differ diff --git a/checkpoint/orbax/checkpoint/experimental/v1/_src/testing/compatibility/managed_checkpoints/v1_managed/no_metrics/no_custom_metadata/0/state/ocdbt.process_0/d/b7a108920476b8a00622e848c3fd6a60 b/checkpoint/orbax/checkpoint/experimental/v1/_src/testing/compatibility/managed_checkpoints/v1_managed/no_metrics/no_custom_metadata/0/state/ocdbt.process_0/d/b7a108920476b8a00622e848c3fd6a60 new file mode 100644 index 000000000..d610a0fd3 Binary files /dev/null and b/checkpoint/orbax/checkpoint/experimental/v1/_src/testing/compatibility/managed_checkpoints/v1_managed/no_metrics/no_custom_metadata/0/state/ocdbt.process_0/d/b7a108920476b8a00622e848c3fd6a60 differ diff --git a/checkpoint/orbax/checkpoint/experimental/v1/_src/testing/compatibility/managed_checkpoints/v1_managed/no_metrics/no_custom_metadata/0/state/ocdbt.process_0/d/bea69c03ab9107699669c614b669908e b/checkpoint/orbax/checkpoint/experimental/v1/_src/testing/compatibility/managed_checkpoints/v1_managed/no_metrics/no_custom_metadata/0/state/ocdbt.process_0/d/bea69c03ab9107699669c614b669908e new file mode 100644 index 000000000..dc3b09d1b Binary files /dev/null and b/checkpoint/orbax/checkpoint/experimental/v1/_src/testing/compatibility/managed_checkpoints/v1_managed/no_metrics/no_custom_metadata/0/state/ocdbt.process_0/d/bea69c03ab9107699669c614b669908e differ diff --git a/checkpoint/orbax/checkpoint/experimental/v1/_src/testing/compatibility/managed_checkpoints/v1_managed/no_metrics/no_custom_metadata/0/state/ocdbt.process_0/manifest.ocdbt b/checkpoint/orbax/checkpoint/experimental/v1/_src/testing/compatibility/managed_checkpoints/v1_managed/no_metrics/no_custom_metadata/0/state/ocdbt.process_0/manifest.ocdbt new file mode 100644 index 000000000..3f96c21d1 Binary files /dev/null and b/checkpoint/orbax/checkpoint/experimental/v1/_src/testing/compatibility/managed_checkpoints/v1_managed/no_metrics/no_custom_metadata/0/state/ocdbt.process_0/manifest.ocdbt differ diff --git a/checkpoint/orbax/checkpoint/experimental/v1/_src/testing/compatibility/pytree_metadata_compatibility_test.py b/checkpoint/orbax/checkpoint/experimental/v1/_src/testing/compatibility/pytree_metadata_compatibility_test_base.py similarity index 81% rename from checkpoint/orbax/checkpoint/experimental/v1/_src/testing/compatibility/pytree_metadata_compatibility_test.py rename to checkpoint/orbax/checkpoint/experimental/v1/_src/testing/compatibility/pytree_metadata_compatibility_test_base.py index 6cfa1f4f4..986c4633e 100644 --- a/checkpoint/orbax/checkpoint/experimental/v1/_src/testing/compatibility/pytree_metadata_compatibility_test.py +++ b/checkpoint/orbax/checkpoint/experimental/v1/_src/testing/compatibility/pytree_metadata_compatibility_test_base.py @@ -27,9 +27,9 @@ from orbax.checkpoint.experimental.v1._src.context import options as options_lib from orbax.checkpoint.experimental.v1._src.handlers import registration from orbax.checkpoint.experimental.v1._src.layout import checkpoint_layout as checkpoint_layout_lib +from orbax.checkpoint.experimental.v1._src.synchronization import multihost from orbax.checkpoint.experimental.v1._src.testing.compatibility import test_utils as compatibility_test_utils - CheckpointLayoutEnum = options_lib.CheckpointLayout InvalidLayoutError = checkpoint_layout_lib.InvalidLayoutError @@ -37,7 +37,8 @@ _BASE_DIR = os.path.join(os.path.dirname(__file__), 'checkpoints') -class PytreeMetadataCompatibilityTest(parameterized.TestCase): +class PytreeMetadataCompatibilityTestBase(parameterized.TestCase): + """Tests for V1 pytree_metadata API against generated Checkpoints.""" def setUp(self) -> None: super().setUp() @@ -155,6 +156,17 @@ def test_pytree_metadata_compatibility( is_pytree: bool, handler_registered: bool, ) -> None: + """Tests pytree_metadata compatibility across V0 and V1 checkpoints. + + Args: + version: The checkpoint version to test against. + checkpointable_name: The name of the checkpointable to load. + name_registered: Whether the checkpointable name is registered. + metadata_present: Whether the checkpoint metadata file is present. + is_direct_checkpoint: Whether the checkpoint is a direct checkpoint. + is_pytree: Whether the checkpointable is a pytree. + handler_registered: Whether the handler is registered. + """ path = compatibility_test_utils.get_checkpoint_path( version, metadata_present, is_direct_checkpoint, is_pytree ) @@ -184,9 +196,12 @@ def test_pytree_metadata_compatibility( path, checkpointable_name=checkpointable_name, ) - test_utils.assert_tree_equal( - self, self.expected_state_metadata, loaded.metadata - ) + expected = self.expected_state_metadata + actual = loaded.metadata + if multihost.is_pathways_backend() or jax.process_count() > 1: + expected = compatibility_test_utils.strip_sharding_metadata(expected) + actual = compatibility_test_utils.strip_sharding_metadata(actual) + test_utils.assert_tree_equal(self, expected, actual) else: with self.assertRaisesRegex(error_type, error_msg): ocp.pytree_metadata( @@ -209,6 +224,12 @@ def test_pytree_metadata_compatibility( def test_pytree_metadata_non_critical_corruptions( self, version: str, alteration: str ) -> None: + """Tests pytree_metadata with non-critical corruptions. + + Args: + version: The checkpoint version to test against. + alteration: The alteration to apply to the checkpoint. + """ path = self.base_dir.joinpath( f'{version}_checkpoints', 'composite_checkpoint', @@ -216,9 +237,12 @@ def test_pytree_metadata_non_critical_corruptions( alteration, ) loaded = ocp.pytree_metadata(path, checkpointable_name='state') - test_utils.assert_tree_equal( - self, self.expected_state_metadata, loaded.metadata - ) + expected = self.expected_state_metadata + actual = loaded.metadata + if multihost.is_pathways_backend() or jax.process_count() > 1: + expected = compatibility_test_utils.strip_sharding_metadata(expected) + actual = compatibility_test_utils.strip_sharding_metadata(actual) + test_utils.assert_tree_equal(self, expected, actual) @parameterized.product( version=['v0', 'v1'], @@ -226,6 +250,11 @@ def test_pytree_metadata_non_critical_corruptions( def test_pytree_metadata_missing_sharding_corruption( self, version: str ) -> None: + """Tests pytree_metadata with missing sharding corruption. + + Args: + version: The checkpoint version to test against. + """ path = self.base_dir.joinpath( f'{version}_checkpoints', 'composite_checkpoint', @@ -247,6 +276,12 @@ def test_pytree_metadata_missing_sharding_corruption( def test_pytree_metadata_critical_corruptions( self, version: str, alteration: str ) -> None: + """Tests pytree_metadata with critical corruptions. + + Args: + version: The checkpoint version to test against. + alteration: The alteration to apply to the checkpoint. + """ path = self.base_dir.joinpath( f'{version}_checkpoints', 'composite_checkpoint', diff --git a/checkpoint/orbax/checkpoint/experimental/v1/_src/testing/compatibility/test_utils.py b/checkpoint/orbax/checkpoint/experimental/v1/_src/testing/compatibility/test_utils.py index 9b3c7eba6..80d2d4581 100644 --- a/checkpoint/orbax/checkpoint/experimental/v1/_src/testing/compatibility/test_utils.py +++ b/checkpoint/orbax/checkpoint/experimental/v1/_src/testing/compatibility/test_utils.py @@ -82,3 +82,31 @@ def create_value_metadata(value: Any) -> Any: return 0.0 else: raise TypeError(f'Unsupported type: {type(value)}') + + +def strip_sharding_metadata(tree: Any) -> Any: + """Strips concrete sharding_metadata from Metadata to decouple from topologies.""" + def _strip(x): + return array_leaf_handler.ArrayMetadata( + shape=x.shape, + dtype=x.dtype, + sharding_metadata=None, + storage_metadata=x.storage_metadata, + ) + return jax.tree.map( + _strip, + tree, + is_leaf=lambda leaf: hasattr(leaf, 'sharding_metadata'), + ) + + +def replicate_on_mesh(tree: Any) -> Any: + """Replicates a PyTree of arrays across all devices in the current mesh.""" + mesh = jax.sharding.Mesh(np.asarray(jax.devices()), ('devices',)) + sharding = jax.sharding.NamedSharding(mesh, jax.sharding.PartitionSpec()) + return jax.tree.map( + lambda x: jax.device_put(x, sharding) + if isinstance(x, (jax.Array, np.ndarray)) + else x, + tree, + ) diff --git a/checkpoint/orbax/checkpoint/experimental/v1/_src/testing/v0v1_compatibility_save_load_test_base.py b/checkpoint/orbax/checkpoint/experimental/v1/_src/testing/v0v1_compatibility_save_load_test_base.py deleted file mode 100644 index 6ea0bfb66..000000000 --- a/checkpoint/orbax/checkpoint/experimental/v1/_src/testing/v0v1_compatibility_save_load_test_base.py +++ /dev/null @@ -1,288 +0,0 @@ -# Copyright 2026 The Orbax Authors. -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. - -"""Base class for v0/v1 compatibility save/load tests.""" - -# pylint: disable=missing-class-docstring,protected-access,missing-function-docstring - -from __future__ import annotations - -from absl.testing import parameterized -from etils import epath -from orbax.checkpoint import args as args_lib -from orbax.checkpoint import test_utils -from orbax.checkpoint._src.checkpointers import checkpointer as v0_checkpointer -from orbax.checkpoint._src.checkpointers import standard_checkpointer -from orbax.checkpoint._src.handlers import composite_checkpoint_handler -import orbax.checkpoint.experimental.v1 as ocp -from orbax.checkpoint.experimental.v1._src.layout import checkpoint_layout -from orbax.checkpoint.experimental.v1._src.path import types as path_types -from orbax.checkpoint.experimental.v1._src.synchronization import multihost -from orbax.checkpoint.experimental.v1._src.testing import array_utils as array_test_utils -from orbax.checkpoint.experimental.v1._src.tree import types as tree_types - -PyTree = tree_types.PyTree -Path = path_types.Path -InvalidLayoutError = checkpoint_layout.InvalidLayoutError - -create_sharded_pytree = array_test_utils.create_sharded_pytree - - -class CompatibilitySaveLoadTestBase: - - class Test(parameterized.TestCase): - - def setUp(self): - super().setUp() - - self.root_directory = epath.Path( - self.create_tempdir(name='root').full_path - ) - self.ckpt_directory = ( - epath.Path(self.create_tempdir(name='direct').full_path) / 'ckpt' - ) - self.pytree, self.abstract_pytree = create_sharded_pytree() - - test_utils.set_tensorstore_driver_for_test() - test_utils.sync_global_processes('CompatibilityTest:setup_complete') - - def tearDown(self): - super().tearDown() - test_utils.sync_global_processes('CompatibilityTest:teardown_complete') - - def save_v0_checkpoint(self, directory: Path): - with standard_checkpointer.StandardCheckpointer() as checkpointer: - checkpointer.save(directory, self.pytree) - - def save_v0_checkpoints( - self, base_dir: Path, *, checkpointable_names: list[str] - ): - args = args_lib.Composite(**{ - checkpointable_name: args_lib.StandardSave(self.pytree) - for checkpointable_name in checkpointable_names - }) - with v0_checkpointer.Checkpointer( - composite_checkpoint_handler.CompositeCheckpointHandler() - ) as checkpointer: - checkpointer.save(base_dir, args) - - def test_async_load(self): - with self.assertRaises(NotImplementedError): - ocp.load_pytree_async( - self.root_directory, - ) - with self.assertRaises(NotImplementedError): - ocp.load_checkpointables_async( - self.root_directory, - ) - - @parameterized.product( - with_abstract_pytree=[True, False], - ) - def test_load_v0_checkpoint_with_v1_load_pytree( - self, with_abstract_pytree: bool - ): - - checkpointable_names = ['default', 'state', 'pytree'] - step_dir = self.root_directory / 'load_pytree_0' - self.save_v0_checkpoints( - step_dir, - checkpointable_names=checkpointable_names, - ) - self.save_v0_checkpoint(self.ckpt_directory) - - with self.subTest('no_checkpointable_name'): - loaded = ocp.load_pytree( - step_dir, - self.abstract_pytree if with_abstract_pytree else None, - ) - test_utils.assert_tree_equal(self, self.pytree, loaded) - - with self.subTest('flat_layout_no_checkpointable_name'): - loaded = ocp.load_pytree( - self.ckpt_directory, - self.abstract_pytree if with_abstract_pytree else None, - ) - test_utils.assert_tree_equal(self, self.pytree, loaded) - - with self.subTest('root_path_no_checkpointable_name_error'): - with self.assertRaises(InvalidLayoutError): - ocp.load_pytree( - self.root_directory, - self.abstract_pytree if with_abstract_pytree else None, - ) - - for checkpointable_name in checkpointable_names: - with self.subTest(f'pass_{checkpointable_name}'): - loaded = ocp.load_pytree( - step_dir, - self.abstract_pytree if with_abstract_pytree else None, - checkpointable_name=checkpointable_name, - ) - test_utils.assert_tree_equal(self, self.pytree, loaded) - - with self.subTest(f'pass_{checkpointable_name}_error'): - with self.assertRaises(InvalidLayoutError): - ocp.load_pytree( - self.ckpt_directory, - self.abstract_pytree if with_abstract_pytree else None, - checkpointable_name=checkpointable_name, - ) - with self.assertRaises(InvalidLayoutError): - ocp.load_pytree( - self.root_directory, - self.abstract_pytree if with_abstract_pytree else None, - checkpointable_name=checkpointable_name, - ) - - with self.subTest('pass_none'): - loaded = ocp.load_pytree( - self.ckpt_directory, - self.abstract_pytree if with_abstract_pytree else None, - checkpointable_name=None, - ) - test_utils.assert_tree_equal(self, self.pytree, loaded) - with self.assertRaises(InvalidLayoutError): - ocp.load_pytree( - step_dir, - self.abstract_pytree if with_abstract_pytree else None, - checkpointable_name=None, - ) - with self.assertRaises(InvalidLayoutError): - ocp.load_pytree( - self.root_directory, - self.abstract_pytree if with_abstract_pytree else None, - checkpointable_name=None, - ) - - def test_load_v0_checkpoint_with_v1_pytree_metadata(self): - checkpointable_names = ['default', 'state', 'pytree'] - step_dir = self.root_directory / 'load_pytree_0' - self.save_v0_checkpoints( - step_dir, - checkpointable_names=checkpointable_names, - ) - self.save_v0_checkpoint(self.ckpt_directory) - - with self.subTest('no_checkpointable_name'): - loaded = ocp.pytree_metadata(step_dir) - test_utils.assert_tree_same_structure( - self, self.abstract_pytree, loaded.metadata - ) - - with self.subTest('flat_layout_no_checkpointable_name'): - metadata = ocp.pytree_metadata(self.ckpt_directory) - test_utils.assert_tree_same_structure( - self, self.abstract_pytree, metadata.metadata - ) - - with self.subTest('root_path_no_checkpointable_name_error'): - with self.assertRaises(InvalidLayoutError): - ocp.pytree_metadata(self.root_directory) - - for checkpointable_name in checkpointable_names: - with self.subTest(f'pass_{checkpointable_name}'): - loaded = ocp.pytree_metadata( - step_dir, - checkpointable_name=checkpointable_name, - ) - test_utils.assert_tree_same_structure( - self, self.abstract_pytree, loaded.metadata - ) - - with self.subTest(f'pass_{checkpointable_name}_error'): - with self.assertRaises(InvalidLayoutError): - ocp.pytree_metadata( - self.ckpt_directory, - checkpointable_name=checkpointable_name, - ) - with self.assertRaises(InvalidLayoutError): - ocp.pytree_metadata( - self.root_directory, - checkpointable_name=checkpointable_name, - ) - - with self.subTest('pass_none'): - loaded = ocp.pytree_metadata( - self.ckpt_directory, - checkpointable_name=None, - ) - test_utils.assert_tree_same_structure( - self, self.abstract_pytree, loaded.metadata - ) - with self.assertRaises(InvalidLayoutError): - ocp.pytree_metadata( - step_dir, - checkpointable_name=None, - ) - with self.assertRaises(InvalidLayoutError): - ocp.pytree_metadata( - self.root_directory, - checkpointable_name=None, - ) - - @parameterized.product( - checkpointable_name=['default', 'state'], - with_abstract_pytree=[True, False], - ) - def test_load_v0_checkpoint_with_v1_load_checkpointables( - self, - checkpointable_name: str, - with_abstract_pytree: bool, - ): - - checkpointable_names = ['default', 'state'] - step_dir = self.root_directory / 'load_checkpointables_0' - self.save_v0_checkpoints( - step_dir, - checkpointable_names=checkpointable_names, - ) - self.save_v0_checkpoint(self.ckpt_directory) - - abstract_checkpointables = ( - {checkpointable_name: self.abstract_pytree} - ) - - with self.subTest('with_context'): - checkpointables_options = ( - ocp.options.CheckpointablesOptions.create_with_handlers( - **{checkpointable_name: ocp.handlers.PyTreeHandler} - ) - ) - with ocp.Context(checkpointables_options=checkpointables_options): - loaded = ocp.load_checkpointables(step_dir, abstract_checkpointables) - test_utils.assert_tree_equal( - self, self.pytree, loaded[checkpointable_name] - ) - - with self.subTest('without_context'): - loaded = ocp.load_checkpointables(step_dir, abstract_checkpointables) - test_utils.assert_tree_equal( - self, self.pytree, loaded[checkpointable_name] - ) - # TODO(b/484400394): Find a better way to inform the user that they need - # to use load_pytree(..., checkpointable_name=None) when item_handlers is - # a str. - with self.subTest('error_with_checkpoint_path'): - with self.assertRaisesRegex( - KeyError, 'Requested checkpointables:' - ): - ocp.load_checkpointables( - self.ckpt_directory, abstract_checkpointables - ) - with self.subTest('error_with_root_path'): - with self.assertRaises(InvalidLayoutError): - ocp.load_checkpointables( - self.root_directory, abstract_checkpointables - ) diff --git a/checkpoint/orbax/checkpoint/experimental/v1/_src/testing/v1_compatibility_load_test.py b/checkpoint/orbax/checkpoint/experimental/v1/_src/testing/v1_compatibility_load_test.py new file mode 100644 index 000000000..b4c2409e8 --- /dev/null +++ b/checkpoint/orbax/checkpoint/experimental/v1/_src/testing/v1_compatibility_load_test.py @@ -0,0 +1,59 @@ +# Copyright 2026 The Orbax Authors. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +from absl import flags +from absl.testing import absltest +from absl.testing import parameterized +import jax +from orbax.checkpoint.experimental.v1._src.testing.compatibility import checkpointables_metadata_compatibility_test_base +from orbax.checkpoint.experimental.v1._src.testing.compatibility import load_checkpointables_compatibility_test_base +from orbax.checkpoint.experimental.v1._src.testing.compatibility import load_pytree_compatibility_test_base +from orbax.checkpoint.experimental.v1._src.testing.compatibility import pytree_metadata_compatibility_test_base + + +FLAGS = flags.FLAGS + +jax.config.update('jax_enable_x64', True) + + +class CheckpointablesMetadataTest( + checkpointables_metadata_compatibility_test_base.CheckpointablesMetadataCompatibilityTestBase, + parameterized.TestCase, +): + pass + + +class LoadCheckpointablesTest( + load_checkpointables_compatibility_test_base.LoadCheckpointablesCompatibilityTestBase, + parameterized.TestCase, +): + pass + + +class LoadPytreeTest( + load_pytree_compatibility_test_base.LoadPytreeCompatibilityTestBase, + parameterized.TestCase, +): + pass + + +class PytreeMetadataTest( + pytree_metadata_compatibility_test_base.PytreeMetadataCompatibilityTestBase, + parameterized.TestCase, +): + pass + + +if __name__ == '__main__': + absltest.main()