From 1cd49cd87d4b8a3c7a5df8eacf93c99830879c6c Mon Sep 17 00:00:00 2001 From: Angel Mau Date: Thu, 7 May 2026 00:49:37 -0700 Subject: [PATCH] Generate static checkpoint managed checkpoints using generate_v0_checkpoint_manager_checkpoints and generate_v1_checkpointer_checkpoints scripts using the following commands: PiperOrigin-RevId: 911796514 --- ...ables_metadata_compatibility_test_base.py} | 44 ++- ...erate_v0_checkpoint_manager_checkpoints.py | 145 +++++++++ .../generate_v1_checkpointer_checkpoints.py | 140 +++++++++ ...heckpointables_compatibility_test_base.py} | 47 ++- ...=> load_pytree_compatibility_test_base.py} | 49 ++- .../0/_CHECKPOINT_METADATA | 1 + .../has_custom_metadata/0/baz/foo.txt | 1 + .../has_custom_metadata/0/metadata/metadata | 1 + .../has_custom_metadata/0/metrics/metrics | 1 + .../has_custom_metadata/0/state/_METADATA | 1 + .../has_custom_metadata/0/state/_sharding | 1 + .../0/state/array_metadatas/process_0 | 1 + .../state/d/ea65b6b4d0fbd9cd118027d3440895d5 | Bin 0 -> 240 bytes .../0/state/descriptor/descriptor.pbtxt | 17 ++ .../uuid-566a6305-6aa3-43ee-bfb0-b020084b7a6f | 1 + .../0/state/manifest.ocdbt | Bin 0 -> 113 bytes .../d/072a9af32a8d72943571b7cd43cb57ea | Bin 0 -> 194 bytes .../d/52008c008bd643144f62fe4ee05ddaea | Bin 0 -> 240 bytes .../d/e85e76a71fa324686df57072d434bd1f | Bin 0 -> 233 bytes .../0/state/ocdbt.process_0/manifest.ocdbt | Bin 0 -> 205 bytes .../metadata/_ROOT_METADATA | 1 + .../no_custom_metadata/0/_CHECKPOINT_METADATA | 1 + .../no_custom_metadata/0/baz/foo.txt | 1 + .../no_custom_metadata/0/metadata/metadata | 1 + .../no_custom_metadata/0/metrics/metrics | 1 + .../no_custom_metadata/0/state/_METADATA | 1 + .../no_custom_metadata/0/state/_sharding | 1 + .../0/state/array_metadatas/process_0 | 1 + .../state/d/e8df828f773a05f06eee72e0a1ad0539 | Bin 0 -> 240 bytes .../0/state/descriptor/descriptor.pbtxt | 17 ++ .../uuid-c29218c5-cc67-4e65-9384-3fbb6119ec28 | 1 + .../no_custom_metadata/0/state/manifest.ocdbt | Bin 0 -> 113 bytes .../d/018ab640ddaf1fdb3258b3291f8d342b | Bin 0 -> 194 bytes .../d/db6c229caac63f4d9254071e22ce2f72 | Bin 0 -> 240 bytes .../0/state/ocdbt.process_0/manifest.ocdbt | Bin 0 -> 166 bytes .../0/_CHECKPOINT_METADATA | 1 + .../has_custom_metadata/0/baz/foo.txt | 1 + .../has_custom_metadata/0/metadata/metadata | 1 + .../has_custom_metadata/0/state/_METADATA | 1 + .../has_custom_metadata/0/state/_sharding | 1 + .../0/state/array_metadatas/process_0 | 1 + .../state/d/e4ce8c428a2582cf4cdf06738cc3e428 | Bin 0 -> 240 bytes .../0/state/descriptor/descriptor.pbtxt | 17 ++ .../uuid-c1461d12-0db6-4b00-83c9-aa74486c53e6 | 1 + .../0/state/manifest.ocdbt | Bin 0 -> 113 bytes .../d/2eee718ac1419a5d32aa79f16581e1f5 | Bin 0 -> 194 bytes .../d/8e88ce9c0c365753a7c3ef9e0a99716f | Bin 0 -> 169 bytes .../d/a97a2190f0a7f8a20f1e3ba20ae0761f | Bin 0 -> 240 bytes .../d/db5bdead14946c3b48f3259e8b2c26fa | Bin 0 -> 233 bytes .../0/state/ocdbt.process_0/manifest.ocdbt | Bin 0 -> 233 bytes .../metadata/_ROOT_METADATA | 1 + .../no_custom_metadata/0/_CHECKPOINT_METADATA | 1 + .../no_custom_metadata/0/baz/foo.txt | 1 + .../no_custom_metadata/0/metadata/metadata | 1 + .../no_custom_metadata/0/state/_METADATA | 1 + .../no_custom_metadata/0/state/_sharding | 1 + .../0/state/array_metadatas/process_0 | 1 + .../state/d/b0f28f119de0006d0f365a3463cca553 | Bin 0 -> 240 bytes .../0/state/descriptor/descriptor.pbtxt | 17 ++ .../uuid-dd7e3917-b5bd-43ae-9aa0-b1af1b48694f | 1 + .../no_custom_metadata/0/state/manifest.ocdbt | Bin 0 -> 113 bytes .../d/16fec959e6da1e856f1c06ebe52298f8 | Bin 0 -> 221 bytes .../d/3aa68a447e9b8b913343fa39c56b6618 | Bin 0 -> 194 bytes .../d/c900507fa78eb4adfa5f14ac0d0de3ec | Bin 0 -> 173 bytes .../d/f74f53bddc8092376c216e00bd1abeb6 | Bin 0 -> 240 bytes .../0/state/ocdbt.process_0/manifest.ocdbt | Bin 0 -> 237 bytes .../0/_CHECKPOINT_METADATA | 1 + .../has_custom_metadata/0/baz/foo.txt | 1 + .../0/descriptor/descriptor.pbtxt | 17 ++ .../uuid-8e9fc4b4-0a03-4f12-b342-e5364a856703 | 1 + .../has_custom_metadata/0/metadata/data.json | 1 + .../has_custom_metadata/0/metrics/metrics | 1 + .../has_custom_metadata/0/state/_METADATA | 1 + .../has_custom_metadata/0/state/_sharding | 1 + .../0/state/array_metadatas/process_0 | 1 + .../state/d/480e6bd7514fd2ae2f11da6a51437c28 | Bin 0 -> 377 bytes .../0/state/manifest.ocdbt | Bin 0 -> 113 bytes .../d/8e18a0958239c8ecf1895b06b1318a01 | Bin 0 -> 343 bytes .../d/a2427bf1070aa986a1826fb23a78654f | Bin 0 -> 377 bytes .../d/f86fde568f5080ee1c2b2adfe13562a8 | Bin 0 -> 293 bytes .../0/state/ocdbt.process_0/manifest.ocdbt | Bin 0 -> 205 bytes .../metadata/_ROOT_METADATA | 1 + .../no_custom_metadata/0/_CHECKPOINT_METADATA | 1 + .../no_custom_metadata/0/baz/foo.txt | 1 + .../0/descriptor/descriptor.pbtxt | 17 ++ .../uuid-508df3ba-9bee-4777-bcfc-e8054d780976 | 1 + .../no_custom_metadata/0/metadata/data.json | 1 + .../no_custom_metadata/0/metrics/metrics | 1 + .../no_custom_metadata/0/state/_METADATA | 1 + .../no_custom_metadata/0/state/_sharding | 1 + .../0/state/array_metadatas/process_0 | 1 + .../state/d/8d9a25bc839918b4fcf92cf52a7124c9 | Bin 0 -> 377 bytes .../no_custom_metadata/0/state/manifest.ocdbt | Bin 0 -> 113 bytes .../d/1494480e45d4ee3eaedc0bd7737d5ab4 | Bin 0 -> 293 bytes .../d/cb04bc578eee9b0493e04285e07c2369 | Bin 0 -> 343 bytes .../d/d0b0a0fcf359d15408bf78f3b2b62945 | Bin 0 -> 377 bytes .../0/state/ocdbt.process_0/manifest.ocdbt | Bin 0 -> 204 bytes .../0/_CHECKPOINT_METADATA | 1 + .../has_custom_metadata/0/baz/foo.txt | 1 + .../0/descriptor/descriptor.pbtxt | 17 ++ .../uuid-c3e829de-2ae2-41d2-9a5b-2120c3777b43 | 1 + .../has_custom_metadata/0/metadata/data.json | 1 + .../has_custom_metadata/0/state/_METADATA | 1 + .../has_custom_metadata/0/state/_sharding | 1 + .../0/state/array_metadatas/process_0 | 1 + .../state/d/ac931c24315ca0ed0940335b3d254e78 | Bin 0 -> 377 bytes .../0/state/manifest.ocdbt | Bin 0 -> 113 bytes .../d/1bdceb1a55e3b70569527962f465bcb3 | Bin 0 -> 293 bytes .../d/5139eb6383ba81d772163c584015116f | Bin 0 -> 377 bytes .../0/state/ocdbt.process_0/manifest.ocdbt | Bin 0 -> 166 bytes .../metadata/_ROOT_METADATA | 1 + .../no_custom_metadata/0/_CHECKPOINT_METADATA | 1 + .../no_custom_metadata/0/baz/foo.txt | 1 + .../0/descriptor/descriptor.pbtxt | 17 ++ .../uuid-3d7ed91d-8849-4c13-804f-7c278ae22f4e | 1 + .../no_custom_metadata/0/metadata/data.json | 1 + .../no_custom_metadata/0/state/_METADATA | 1 + .../no_custom_metadata/0/state/_sharding | 1 + .../0/state/array_metadatas/process_0 | 1 + .../state/d/d759cea92eb7825e071cdaf236421308 | Bin 0 -> 377 bytes .../no_custom_metadata/0/state/manifest.ocdbt | Bin 0 -> 113 bytes .../d/b47b7183ae4c1865684f50c448253b5a | Bin 0 -> 343 bytes .../d/b7a108920476b8a00622e848c3fd6a60 | Bin 0 -> 293 bytes .../d/bea69c03ab9107699669c614b669908e | Bin 0 -> 377 bytes .../0/state/ocdbt.process_0/manifest.ocdbt | Bin 0 -> 199 bytes ...ytree_metadata_compatibility_test_base.py} | 51 +++- .../_src/testing/compatibility/test_utils.py | 28 ++ .../v0v1_compatibility_save_load_test_base.py | 288 ------------------ .../testing/v1_compatibility_load_test.py | 59 ++++ 129 files changed, 747 insertions(+), 304 deletions(-) rename checkpoint/orbax/checkpoint/experimental/v1/_src/testing/compatibility/{checkpointables_metadata_compatibility_test.py => checkpointables_metadata_compatibility_test_base.py} (80%) create mode 100644 checkpoint/orbax/checkpoint/experimental/v1/_src/testing/compatibility/generate_v0_checkpoint_manager_checkpoints.py create mode 100644 checkpoint/orbax/checkpoint/experimental/v1/_src/testing/compatibility/generate_v1_checkpointer_checkpoints.py rename checkpoint/orbax/checkpoint/experimental/v1/_src/testing/compatibility/{load_checkpointables_compatibility_test.py => load_checkpointables_compatibility_test_base.py} (83%) rename checkpoint/orbax/checkpoint/experimental/v1/_src/testing/compatibility/{load_pytree_compatibility_test.py => load_pytree_compatibility_test_base.py} (83%) create mode 100644 checkpoint/orbax/checkpoint/experimental/v1/_src/testing/compatibility/managed_checkpoints/v0_managed/has_metrics/has_custom_metadata/0/_CHECKPOINT_METADATA create mode 100644 checkpoint/orbax/checkpoint/experimental/v1/_src/testing/compatibility/managed_checkpoints/v0_managed/has_metrics/has_custom_metadata/0/baz/foo.txt create mode 100644 checkpoint/orbax/checkpoint/experimental/v1/_src/testing/compatibility/managed_checkpoints/v0_managed/has_metrics/has_custom_metadata/0/metadata/metadata create mode 100644 checkpoint/orbax/checkpoint/experimental/v1/_src/testing/compatibility/managed_checkpoints/v0_managed/has_metrics/has_custom_metadata/0/metrics/metrics create mode 100644 checkpoint/orbax/checkpoint/experimental/v1/_src/testing/compatibility/managed_checkpoints/v0_managed/has_metrics/has_custom_metadata/0/state/_METADATA create mode 100644 checkpoint/orbax/checkpoint/experimental/v1/_src/testing/compatibility/managed_checkpoints/v0_managed/has_metrics/has_custom_metadata/0/state/_sharding create mode 100644 checkpoint/orbax/checkpoint/experimental/v1/_src/testing/compatibility/managed_checkpoints/v0_managed/has_metrics/has_custom_metadata/0/state/array_metadatas/process_0 create mode 100644 checkpoint/orbax/checkpoint/experimental/v1/_src/testing/compatibility/managed_checkpoints/v0_managed/has_metrics/has_custom_metadata/0/state/d/ea65b6b4d0fbd9cd118027d3440895d5 create mode 100644 checkpoint/orbax/checkpoint/experimental/v1/_src/testing/compatibility/managed_checkpoints/v0_managed/has_metrics/has_custom_metadata/0/state/descriptor/descriptor.pbtxt create mode 100644 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 create mode 100644 checkpoint/orbax/checkpoint/experimental/v1/_src/testing/compatibility/managed_checkpoints/v0_managed/has_metrics/has_custom_metadata/0/state/manifest.ocdbt create mode 100644 checkpoint/orbax/checkpoint/experimental/v1/_src/testing/compatibility/managed_checkpoints/v0_managed/has_metrics/has_custom_metadata/0/state/ocdbt.process_0/d/072a9af32a8d72943571b7cd43cb57ea create mode 100644 checkpoint/orbax/checkpoint/experimental/v1/_src/testing/compatibility/managed_checkpoints/v0_managed/has_metrics/has_custom_metadata/0/state/ocdbt.process_0/d/52008c008bd643144f62fe4ee05ddaea create mode 100644 checkpoint/orbax/checkpoint/experimental/v1/_src/testing/compatibility/managed_checkpoints/v0_managed/has_metrics/has_custom_metadata/0/state/ocdbt.process_0/d/e85e76a71fa324686df57072d434bd1f create mode 100644 checkpoint/orbax/checkpoint/experimental/v1/_src/testing/compatibility/managed_checkpoints/v0_managed/has_metrics/has_custom_metadata/0/state/ocdbt.process_0/manifest.ocdbt create mode 100644 checkpoint/orbax/checkpoint/experimental/v1/_src/testing/compatibility/managed_checkpoints/v0_managed/has_metrics/has_custom_metadata/metadata/_ROOT_METADATA create mode 100644 checkpoint/orbax/checkpoint/experimental/v1/_src/testing/compatibility/managed_checkpoints/v0_managed/has_metrics/no_custom_metadata/0/_CHECKPOINT_METADATA create mode 100644 checkpoint/orbax/checkpoint/experimental/v1/_src/testing/compatibility/managed_checkpoints/v0_managed/has_metrics/no_custom_metadata/0/baz/foo.txt create mode 100644 checkpoint/orbax/checkpoint/experimental/v1/_src/testing/compatibility/managed_checkpoints/v0_managed/has_metrics/no_custom_metadata/0/metadata/metadata create mode 100644 checkpoint/orbax/checkpoint/experimental/v1/_src/testing/compatibility/managed_checkpoints/v0_managed/has_metrics/no_custom_metadata/0/metrics/metrics create mode 100644 checkpoint/orbax/checkpoint/experimental/v1/_src/testing/compatibility/managed_checkpoints/v0_managed/has_metrics/no_custom_metadata/0/state/_METADATA create mode 100644 checkpoint/orbax/checkpoint/experimental/v1/_src/testing/compatibility/managed_checkpoints/v0_managed/has_metrics/no_custom_metadata/0/state/_sharding create mode 100644 checkpoint/orbax/checkpoint/experimental/v1/_src/testing/compatibility/managed_checkpoints/v0_managed/has_metrics/no_custom_metadata/0/state/array_metadatas/process_0 create mode 100644 checkpoint/orbax/checkpoint/experimental/v1/_src/testing/compatibility/managed_checkpoints/v0_managed/has_metrics/no_custom_metadata/0/state/d/e8df828f773a05f06eee72e0a1ad0539 create mode 100644 checkpoint/orbax/checkpoint/experimental/v1/_src/testing/compatibility/managed_checkpoints/v0_managed/has_metrics/no_custom_metadata/0/state/descriptor/descriptor.pbtxt create mode 100644 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 create mode 100644 checkpoint/orbax/checkpoint/experimental/v1/_src/testing/compatibility/managed_checkpoints/v0_managed/has_metrics/no_custom_metadata/0/state/manifest.ocdbt create mode 100644 checkpoint/orbax/checkpoint/experimental/v1/_src/testing/compatibility/managed_checkpoints/v0_managed/has_metrics/no_custom_metadata/0/state/ocdbt.process_0/d/018ab640ddaf1fdb3258b3291f8d342b create mode 100644 checkpoint/orbax/checkpoint/experimental/v1/_src/testing/compatibility/managed_checkpoints/v0_managed/has_metrics/no_custom_metadata/0/state/ocdbt.process_0/d/db6c229caac63f4d9254071e22ce2f72 create mode 100644 checkpoint/orbax/checkpoint/experimental/v1/_src/testing/compatibility/managed_checkpoints/v0_managed/has_metrics/no_custom_metadata/0/state/ocdbt.process_0/manifest.ocdbt create mode 100644 checkpoint/orbax/checkpoint/experimental/v1/_src/testing/compatibility/managed_checkpoints/v0_managed/no_metrics/has_custom_metadata/0/_CHECKPOINT_METADATA create mode 100644 checkpoint/orbax/checkpoint/experimental/v1/_src/testing/compatibility/managed_checkpoints/v0_managed/no_metrics/has_custom_metadata/0/baz/foo.txt create mode 100644 checkpoint/orbax/checkpoint/experimental/v1/_src/testing/compatibility/managed_checkpoints/v0_managed/no_metrics/has_custom_metadata/0/metadata/metadata create mode 100644 checkpoint/orbax/checkpoint/experimental/v1/_src/testing/compatibility/managed_checkpoints/v0_managed/no_metrics/has_custom_metadata/0/state/_METADATA create mode 100644 checkpoint/orbax/checkpoint/experimental/v1/_src/testing/compatibility/managed_checkpoints/v0_managed/no_metrics/has_custom_metadata/0/state/_sharding create mode 100644 checkpoint/orbax/checkpoint/experimental/v1/_src/testing/compatibility/managed_checkpoints/v0_managed/no_metrics/has_custom_metadata/0/state/array_metadatas/process_0 create mode 100644 checkpoint/orbax/checkpoint/experimental/v1/_src/testing/compatibility/managed_checkpoints/v0_managed/no_metrics/has_custom_metadata/0/state/d/e4ce8c428a2582cf4cdf06738cc3e428 create mode 100644 checkpoint/orbax/checkpoint/experimental/v1/_src/testing/compatibility/managed_checkpoints/v0_managed/no_metrics/has_custom_metadata/0/state/descriptor/descriptor.pbtxt create mode 100644 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 create mode 100644 checkpoint/orbax/checkpoint/experimental/v1/_src/testing/compatibility/managed_checkpoints/v0_managed/no_metrics/has_custom_metadata/0/state/manifest.ocdbt create mode 100644 checkpoint/orbax/checkpoint/experimental/v1/_src/testing/compatibility/managed_checkpoints/v0_managed/no_metrics/has_custom_metadata/0/state/ocdbt.process_0/d/2eee718ac1419a5d32aa79f16581e1f5 create mode 100644 checkpoint/orbax/checkpoint/experimental/v1/_src/testing/compatibility/managed_checkpoints/v0_managed/no_metrics/has_custom_metadata/0/state/ocdbt.process_0/d/8e88ce9c0c365753a7c3ef9e0a99716f create mode 100644 checkpoint/orbax/checkpoint/experimental/v1/_src/testing/compatibility/managed_checkpoints/v0_managed/no_metrics/has_custom_metadata/0/state/ocdbt.process_0/d/a97a2190f0a7f8a20f1e3ba20ae0761f create mode 100644 checkpoint/orbax/checkpoint/experimental/v1/_src/testing/compatibility/managed_checkpoints/v0_managed/no_metrics/has_custom_metadata/0/state/ocdbt.process_0/d/db5bdead14946c3b48f3259e8b2c26fa create mode 100644 checkpoint/orbax/checkpoint/experimental/v1/_src/testing/compatibility/managed_checkpoints/v0_managed/no_metrics/has_custom_metadata/0/state/ocdbt.process_0/manifest.ocdbt create mode 100644 checkpoint/orbax/checkpoint/experimental/v1/_src/testing/compatibility/managed_checkpoints/v0_managed/no_metrics/has_custom_metadata/metadata/_ROOT_METADATA create mode 100644 checkpoint/orbax/checkpoint/experimental/v1/_src/testing/compatibility/managed_checkpoints/v0_managed/no_metrics/no_custom_metadata/0/_CHECKPOINT_METADATA create mode 100644 checkpoint/orbax/checkpoint/experimental/v1/_src/testing/compatibility/managed_checkpoints/v0_managed/no_metrics/no_custom_metadata/0/baz/foo.txt create mode 100644 checkpoint/orbax/checkpoint/experimental/v1/_src/testing/compatibility/managed_checkpoints/v0_managed/no_metrics/no_custom_metadata/0/metadata/metadata create mode 100644 checkpoint/orbax/checkpoint/experimental/v1/_src/testing/compatibility/managed_checkpoints/v0_managed/no_metrics/no_custom_metadata/0/state/_METADATA create mode 100644 checkpoint/orbax/checkpoint/experimental/v1/_src/testing/compatibility/managed_checkpoints/v0_managed/no_metrics/no_custom_metadata/0/state/_sharding create mode 100644 checkpoint/orbax/checkpoint/experimental/v1/_src/testing/compatibility/managed_checkpoints/v0_managed/no_metrics/no_custom_metadata/0/state/array_metadatas/process_0 create mode 100644 checkpoint/orbax/checkpoint/experimental/v1/_src/testing/compatibility/managed_checkpoints/v0_managed/no_metrics/no_custom_metadata/0/state/d/b0f28f119de0006d0f365a3463cca553 create mode 100644 checkpoint/orbax/checkpoint/experimental/v1/_src/testing/compatibility/managed_checkpoints/v0_managed/no_metrics/no_custom_metadata/0/state/descriptor/descriptor.pbtxt create mode 100644 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 create mode 100644 checkpoint/orbax/checkpoint/experimental/v1/_src/testing/compatibility/managed_checkpoints/v0_managed/no_metrics/no_custom_metadata/0/state/manifest.ocdbt create mode 100644 checkpoint/orbax/checkpoint/experimental/v1/_src/testing/compatibility/managed_checkpoints/v0_managed/no_metrics/no_custom_metadata/0/state/ocdbt.process_0/d/16fec959e6da1e856f1c06ebe52298f8 create mode 100644 checkpoint/orbax/checkpoint/experimental/v1/_src/testing/compatibility/managed_checkpoints/v0_managed/no_metrics/no_custom_metadata/0/state/ocdbt.process_0/d/3aa68a447e9b8b913343fa39c56b6618 create mode 100644 checkpoint/orbax/checkpoint/experimental/v1/_src/testing/compatibility/managed_checkpoints/v0_managed/no_metrics/no_custom_metadata/0/state/ocdbt.process_0/d/c900507fa78eb4adfa5f14ac0d0de3ec create mode 100644 checkpoint/orbax/checkpoint/experimental/v1/_src/testing/compatibility/managed_checkpoints/v0_managed/no_metrics/no_custom_metadata/0/state/ocdbt.process_0/d/f74f53bddc8092376c216e00bd1abeb6 create mode 100644 checkpoint/orbax/checkpoint/experimental/v1/_src/testing/compatibility/managed_checkpoints/v0_managed/no_metrics/no_custom_metadata/0/state/ocdbt.process_0/manifest.ocdbt create mode 100644 checkpoint/orbax/checkpoint/experimental/v1/_src/testing/compatibility/managed_checkpoints/v1_managed/has_metrics/has_custom_metadata/0/_CHECKPOINT_METADATA create mode 100644 checkpoint/orbax/checkpoint/experimental/v1/_src/testing/compatibility/managed_checkpoints/v1_managed/has_metrics/has_custom_metadata/0/baz/foo.txt create mode 100644 checkpoint/orbax/checkpoint/experimental/v1/_src/testing/compatibility/managed_checkpoints/v1_managed/has_metrics/has_custom_metadata/0/descriptor/descriptor.pbtxt create mode 100644 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 create mode 100644 checkpoint/orbax/checkpoint/experimental/v1/_src/testing/compatibility/managed_checkpoints/v1_managed/has_metrics/has_custom_metadata/0/metadata/data.json create mode 100644 checkpoint/orbax/checkpoint/experimental/v1/_src/testing/compatibility/managed_checkpoints/v1_managed/has_metrics/has_custom_metadata/0/metrics/metrics create mode 100644 checkpoint/orbax/checkpoint/experimental/v1/_src/testing/compatibility/managed_checkpoints/v1_managed/has_metrics/has_custom_metadata/0/state/_METADATA create mode 100644 checkpoint/orbax/checkpoint/experimental/v1/_src/testing/compatibility/managed_checkpoints/v1_managed/has_metrics/has_custom_metadata/0/state/_sharding create mode 100644 checkpoint/orbax/checkpoint/experimental/v1/_src/testing/compatibility/managed_checkpoints/v1_managed/has_metrics/has_custom_metadata/0/state/array_metadatas/process_0 create mode 100644 checkpoint/orbax/checkpoint/experimental/v1/_src/testing/compatibility/managed_checkpoints/v1_managed/has_metrics/has_custom_metadata/0/state/d/480e6bd7514fd2ae2f11da6a51437c28 create mode 100644 checkpoint/orbax/checkpoint/experimental/v1/_src/testing/compatibility/managed_checkpoints/v1_managed/has_metrics/has_custom_metadata/0/state/manifest.ocdbt create mode 100644 checkpoint/orbax/checkpoint/experimental/v1/_src/testing/compatibility/managed_checkpoints/v1_managed/has_metrics/has_custom_metadata/0/state/ocdbt.process_0/d/8e18a0958239c8ecf1895b06b1318a01 create mode 100644 checkpoint/orbax/checkpoint/experimental/v1/_src/testing/compatibility/managed_checkpoints/v1_managed/has_metrics/has_custom_metadata/0/state/ocdbt.process_0/d/a2427bf1070aa986a1826fb23a78654f create mode 100644 checkpoint/orbax/checkpoint/experimental/v1/_src/testing/compatibility/managed_checkpoints/v1_managed/has_metrics/has_custom_metadata/0/state/ocdbt.process_0/d/f86fde568f5080ee1c2b2adfe13562a8 create mode 100644 checkpoint/orbax/checkpoint/experimental/v1/_src/testing/compatibility/managed_checkpoints/v1_managed/has_metrics/has_custom_metadata/0/state/ocdbt.process_0/manifest.ocdbt create mode 100644 checkpoint/orbax/checkpoint/experimental/v1/_src/testing/compatibility/managed_checkpoints/v1_managed/has_metrics/has_custom_metadata/metadata/_ROOT_METADATA create mode 100644 checkpoint/orbax/checkpoint/experimental/v1/_src/testing/compatibility/managed_checkpoints/v1_managed/has_metrics/no_custom_metadata/0/_CHECKPOINT_METADATA create mode 100644 checkpoint/orbax/checkpoint/experimental/v1/_src/testing/compatibility/managed_checkpoints/v1_managed/has_metrics/no_custom_metadata/0/baz/foo.txt create mode 100644 checkpoint/orbax/checkpoint/experimental/v1/_src/testing/compatibility/managed_checkpoints/v1_managed/has_metrics/no_custom_metadata/0/descriptor/descriptor.pbtxt create mode 100644 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 create mode 100644 checkpoint/orbax/checkpoint/experimental/v1/_src/testing/compatibility/managed_checkpoints/v1_managed/has_metrics/no_custom_metadata/0/metadata/data.json create mode 100644 checkpoint/orbax/checkpoint/experimental/v1/_src/testing/compatibility/managed_checkpoints/v1_managed/has_metrics/no_custom_metadata/0/metrics/metrics create mode 100644 checkpoint/orbax/checkpoint/experimental/v1/_src/testing/compatibility/managed_checkpoints/v1_managed/has_metrics/no_custom_metadata/0/state/_METADATA create mode 100644 checkpoint/orbax/checkpoint/experimental/v1/_src/testing/compatibility/managed_checkpoints/v1_managed/has_metrics/no_custom_metadata/0/state/_sharding create mode 100644 checkpoint/orbax/checkpoint/experimental/v1/_src/testing/compatibility/managed_checkpoints/v1_managed/has_metrics/no_custom_metadata/0/state/array_metadatas/process_0 create mode 100644 checkpoint/orbax/checkpoint/experimental/v1/_src/testing/compatibility/managed_checkpoints/v1_managed/has_metrics/no_custom_metadata/0/state/d/8d9a25bc839918b4fcf92cf52a7124c9 create mode 100644 checkpoint/orbax/checkpoint/experimental/v1/_src/testing/compatibility/managed_checkpoints/v1_managed/has_metrics/no_custom_metadata/0/state/manifest.ocdbt create mode 100644 checkpoint/orbax/checkpoint/experimental/v1/_src/testing/compatibility/managed_checkpoints/v1_managed/has_metrics/no_custom_metadata/0/state/ocdbt.process_0/d/1494480e45d4ee3eaedc0bd7737d5ab4 create mode 100644 checkpoint/orbax/checkpoint/experimental/v1/_src/testing/compatibility/managed_checkpoints/v1_managed/has_metrics/no_custom_metadata/0/state/ocdbt.process_0/d/cb04bc578eee9b0493e04285e07c2369 create mode 100644 checkpoint/orbax/checkpoint/experimental/v1/_src/testing/compatibility/managed_checkpoints/v1_managed/has_metrics/no_custom_metadata/0/state/ocdbt.process_0/d/d0b0a0fcf359d15408bf78f3b2b62945 create mode 100644 checkpoint/orbax/checkpoint/experimental/v1/_src/testing/compatibility/managed_checkpoints/v1_managed/has_metrics/no_custom_metadata/0/state/ocdbt.process_0/manifest.ocdbt create mode 100644 checkpoint/orbax/checkpoint/experimental/v1/_src/testing/compatibility/managed_checkpoints/v1_managed/no_metrics/has_custom_metadata/0/_CHECKPOINT_METADATA create mode 100644 checkpoint/orbax/checkpoint/experimental/v1/_src/testing/compatibility/managed_checkpoints/v1_managed/no_metrics/has_custom_metadata/0/baz/foo.txt create mode 100644 checkpoint/orbax/checkpoint/experimental/v1/_src/testing/compatibility/managed_checkpoints/v1_managed/no_metrics/has_custom_metadata/0/descriptor/descriptor.pbtxt create mode 100644 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 create mode 100644 checkpoint/orbax/checkpoint/experimental/v1/_src/testing/compatibility/managed_checkpoints/v1_managed/no_metrics/has_custom_metadata/0/metadata/data.json create mode 100644 checkpoint/orbax/checkpoint/experimental/v1/_src/testing/compatibility/managed_checkpoints/v1_managed/no_metrics/has_custom_metadata/0/state/_METADATA create mode 100644 checkpoint/orbax/checkpoint/experimental/v1/_src/testing/compatibility/managed_checkpoints/v1_managed/no_metrics/has_custom_metadata/0/state/_sharding create mode 100644 checkpoint/orbax/checkpoint/experimental/v1/_src/testing/compatibility/managed_checkpoints/v1_managed/no_metrics/has_custom_metadata/0/state/array_metadatas/process_0 create mode 100644 checkpoint/orbax/checkpoint/experimental/v1/_src/testing/compatibility/managed_checkpoints/v1_managed/no_metrics/has_custom_metadata/0/state/d/ac931c24315ca0ed0940335b3d254e78 create mode 100644 checkpoint/orbax/checkpoint/experimental/v1/_src/testing/compatibility/managed_checkpoints/v1_managed/no_metrics/has_custom_metadata/0/state/manifest.ocdbt create mode 100644 checkpoint/orbax/checkpoint/experimental/v1/_src/testing/compatibility/managed_checkpoints/v1_managed/no_metrics/has_custom_metadata/0/state/ocdbt.process_0/d/1bdceb1a55e3b70569527962f465bcb3 create mode 100644 checkpoint/orbax/checkpoint/experimental/v1/_src/testing/compatibility/managed_checkpoints/v1_managed/no_metrics/has_custom_metadata/0/state/ocdbt.process_0/d/5139eb6383ba81d772163c584015116f create mode 100644 checkpoint/orbax/checkpoint/experimental/v1/_src/testing/compatibility/managed_checkpoints/v1_managed/no_metrics/has_custom_metadata/0/state/ocdbt.process_0/manifest.ocdbt create mode 100644 checkpoint/orbax/checkpoint/experimental/v1/_src/testing/compatibility/managed_checkpoints/v1_managed/no_metrics/has_custom_metadata/metadata/_ROOT_METADATA create mode 100644 checkpoint/orbax/checkpoint/experimental/v1/_src/testing/compatibility/managed_checkpoints/v1_managed/no_metrics/no_custom_metadata/0/_CHECKPOINT_METADATA create mode 100644 checkpoint/orbax/checkpoint/experimental/v1/_src/testing/compatibility/managed_checkpoints/v1_managed/no_metrics/no_custom_metadata/0/baz/foo.txt create mode 100644 checkpoint/orbax/checkpoint/experimental/v1/_src/testing/compatibility/managed_checkpoints/v1_managed/no_metrics/no_custom_metadata/0/descriptor/descriptor.pbtxt create mode 100644 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 create mode 100644 checkpoint/orbax/checkpoint/experimental/v1/_src/testing/compatibility/managed_checkpoints/v1_managed/no_metrics/no_custom_metadata/0/metadata/data.json create mode 100644 checkpoint/orbax/checkpoint/experimental/v1/_src/testing/compatibility/managed_checkpoints/v1_managed/no_metrics/no_custom_metadata/0/state/_METADATA create mode 100644 checkpoint/orbax/checkpoint/experimental/v1/_src/testing/compatibility/managed_checkpoints/v1_managed/no_metrics/no_custom_metadata/0/state/_sharding create mode 100644 checkpoint/orbax/checkpoint/experimental/v1/_src/testing/compatibility/managed_checkpoints/v1_managed/no_metrics/no_custom_metadata/0/state/array_metadatas/process_0 create mode 100644 checkpoint/orbax/checkpoint/experimental/v1/_src/testing/compatibility/managed_checkpoints/v1_managed/no_metrics/no_custom_metadata/0/state/d/d759cea92eb7825e071cdaf236421308 create mode 100644 checkpoint/orbax/checkpoint/experimental/v1/_src/testing/compatibility/managed_checkpoints/v1_managed/no_metrics/no_custom_metadata/0/state/manifest.ocdbt create mode 100644 checkpoint/orbax/checkpoint/experimental/v1/_src/testing/compatibility/managed_checkpoints/v1_managed/no_metrics/no_custom_metadata/0/state/ocdbt.process_0/d/b47b7183ae4c1865684f50c448253b5a create mode 100644 checkpoint/orbax/checkpoint/experimental/v1/_src/testing/compatibility/managed_checkpoints/v1_managed/no_metrics/no_custom_metadata/0/state/ocdbt.process_0/d/b7a108920476b8a00622e848c3fd6a60 create mode 100644 checkpoint/orbax/checkpoint/experimental/v1/_src/testing/compatibility/managed_checkpoints/v1_managed/no_metrics/no_custom_metadata/0/state/ocdbt.process_0/d/bea69c03ab9107699669c614b669908e create mode 100644 checkpoint/orbax/checkpoint/experimental/v1/_src/testing/compatibility/managed_checkpoints/v1_managed/no_metrics/no_custom_metadata/0/state/ocdbt.process_0/manifest.ocdbt rename checkpoint/orbax/checkpoint/experimental/v1/_src/testing/compatibility/{pytree_metadata_compatibility_test.py => pytree_metadata_compatibility_test_base.py} (81%) delete mode 100644 checkpoint/orbax/checkpoint/experimental/v1/_src/testing/v0v1_compatibility_save_load_test_base.py create mode 100644 checkpoint/orbax/checkpoint/experimental/v1/_src/testing/v1_compatibility_load_test.py 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 0000000000000000000000000000000000000000..57b9216da50c0b7edea3ca17acf439ebc4b6a613 GIT binary patch literal 240 zcmVT{yjbiF q08Aqb08aoE3pH|6EY55rWy=yq0ml2M5v zML#vs%rwa?$t1-fEh)t^ImOV>!obKp#n{Bez{1ir#gq}G^aCTy2Sx_&jp8onE!Ijf JtXUb=1pt(KA&vk5 literal 0 HcmV?d00001 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 0000000000000000000000000000000000000000..8d6d8647aec8bf666fcf3369e4730acaca153dbb GIT binary patch literal 194 zcmV;z06qT<+aTV;00000000000VuUE{a|PSH3a|y${{0gC(Z%<|Ht0R|C>ycOlj52 zHqryp{~{7pVO0Kj(tmtFY$S#u4dkOP43k5Z{`$f5yo_P$gi|Uy?xHt+VdyEH?B!Ymf2_yO)^-B3}sw2>D)kje4!&HbN~PV literal 0 HcmV?d00001 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 0000000000000000000000000000000000000000..57b9216da50c0b7edea3ca17acf439ebc4b6a613 GIT binary patch literal 240 zcmVT{yjbiF q08Aqb08aoE3pH|6EY55rWy=yq0S60{XhQST0&BdhJft=+W@u~wiUJ+d$A`N5(Simd|c2z zhqE5=>J87YvVCwhIyXgmo!o*r&Nau?6&l0yDn%l@dh@t;NeZ3lH5e(>*cAmK?-Zi$ z>MTOZ^Q?W6-v^FOED(14U;WP7zw-zRVL-tyl(v zy*eOe5=pBu5F)uHS9TaIl>{vK7XwM;P63-l@=-9rzz{Sb93+u^04QWg>@KoY!bld$ zp_WRkwx9g?U9*Puqm1mH*VZ1JLP?$aKV+0Ltu9vP_Bj5Oj6ISu2c0&Y&*Oi3wuX-# z*Du|@>W;f+&F}8AFiNf6-Gnp~4k!RxqXb40!AQC7b1BZnK04t%gM)DDYVFy79 HPW$2*!T{yjbiF q08Aqb08aoE3pH|6EY55rWy=yq0iOXttr0-sI@+(6FN5i4(`Yenfm5fRZ zDf+1vDQOl)7HQ_@#)$@|X$EGgsj22hsRoILi75uA#+HmAr5_ksJ}@%m8`Z2lZ?RT_ J;YhwuDFF2KBS-)M literal 0 HcmV?d00001 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 0000000000000000000000000000000000000000..8d6d8647aec8bf666fcf3369e4730acaca153dbb GIT binary patch literal 194 zcmV;z06qT<+aTV;00000000000VuUE{a|PSH3a|y${{0gC(Z%<|Ht0R|C>ycOlj52 zHqryp{~{7pVO0Kj(tmtFY$S#u4dkOP43k5Z{`$f5yo_P$gi|Uy?xHt+VdyEH?B!Ymf2_yO)^-B3}sw2>D)kje4!&HbN~PV literal 0 HcmV?d00001 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 0000000000000000000000000000000000000000..57b9216da50c0b7edea3ca17acf439ebc4b6a613 GIT binary patch literal 240 zcmVT{yjbiF q08Aqb08aoE3pH|6EY55rWy=yq0?=#FGz&|kb?#P-{jmF-P@H@*oQ++t|pXgG9TpM?>ml8H%4 zfq@}K-@wo!G0DutASER+%`hz`$=Jx$0!UaIrdgyIn;0dfB$*`}8CfPLCMKI1rZ@V86~G0rI{NsF){&70~&dV@dG0h3y5T3=uI`Ac-~^I#F6Xi%Rn>(T{yjbiF q08Aqb08aoE3pH|6EY55rWy=yq0ycOlj52 zHqryp{~{7pVO0Kj(tmtFY$S#u4dkOP43k5Z{`$f5yo_P$gi|Uy?xHt+VdyEH?B!Ymf2_yO)^-B3}sw2>D)kje4!&HbN~PV literal 0 HcmV?d00001 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 0000000000000000000000000000000000000000..a8adc2748e584922acd5e8419b77bcfceeb2d04c GIT binary patch literal 169 zcmV;a09OAD+aTVl00000000000VuUE{UEe;1OS4`9v^_mIWsV~YhQmtWVtrpPD?8t zO;mam0|H_@;3yc#w$mUVwc)oU3zg}sAD-u3{F*nm=D7SiO`oJSwp^XS49_ow;<`HW zsH|{1agVjzhOMhfkZ=2JyZVY9J-^tcJiqQ`*gWMbQ9$Hd;T%S}WE=}fnKeQJ3OEcS X1ppxnV8Y}jlgMh`SJ6^tZxAXgGapFg literal 0 HcmV?d00001 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 0000000000000000000000000000000000000000..57b9216da50c0b7edea3ca17acf439ebc4b6a613 GIT binary patch literal 240 zcmVT{yjbiF q08Aqb08aoE3pH|6EY55rWy=yq0S60{XhQST0&BdhJft=+W@u~wiUJ+d$A`N5(Simd|c2z zhqE5=>J87YvVCwhIyXgmo!o*r&Nau?6&l0yDn%l@dh@t;NeZ3lH5e(>*cAmK?-Zi$ z>MTOZ^Q?W6-v^FOED(14U;WP7zw-zRVL-tyl1YooGPp}VujWEiGEDYYPkP+A%h&I90q79NLuP?(S~ jKm;fdAWRYn03gUpBna?AU}x}Ekb}d7?1ST{yjbiF q08Aqb08aoE3pH|6EY55rWy=yq0haBUt$oE=h5madO=D=_XgG9TpM?>ml2M5v zML)?P&B!9n(9kj^)xf~OEX5$r*vvH1*u=~@IXThP)R+;Z^aCTy2S$edwzH33uvja> J;8Sz#Apj8v(PCJ?$W-I@LWo~6$E)=S zew%Z;kkgu{hXn_`i+G|&^N4K%WuiycOlj52 zHqryp{~{7pVO0Kj(tmtFY$S#u4dkOP43k5Z{`$f5yo_P$gi|Uy?xHt+VdyEH?B!Ymf2_yO)^-B3}sw2>D)kje4!&HbN~PV literal 0 HcmV?d00001 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 0000000000000000000000000000000000000000..317f53f948e4b689f605a9edaf542dae29cce0a7 GIT binary patch literal 173 zcmV;e08;-9+aTVp00000000000VuUE{UElL1OTGRA0cox&H?=Y$KJ{Ro6K@;)mJUl zWTOC34q+rP%9aNxF1-B{As@Bjw8m52=Ux1oH@4=u{5nmaq&2o&-M|jdFNET{ zI`XKja6562wcCcRt7(vL`)s@ViXA<_*rhzb?q%3K_jB?310w8792oWgA bB5)W+1ppxnV8Y}jlgMh`U(r%#VJE8kws}!e literal 0 HcmV?d00001 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 0000000000000000000000000000000000000000..57b9216da50c0b7edea3ca17acf439ebc4b6a613 GIT binary patch literal 240 zcmVT{yjbiF q08Aqb08aoE3pH|6EY55rWy=yq0ex_%De+!r9o;${pzOFh$~kQ6L#GkP{LQ2#<#;5dawtsgg9( zBs~jCYAq_Q_4j7dRexC%*J%w@j#7Gi(>XJKd#O@mJ2qcW+0LJxo{m0M7@6(*z6e{- z=2~OW+2K3RHE(z0$*^hO`>#ee_wloq`|!c|a~VW33?rbFS_mPO7UW_Q35cMCjtay& npio#WHWn%vI0ym=03gUpBna?AU}x}Ekb}d7?1Sf*CUjN=wrv9X zWZx6HZCmU{h}PzHj#ipA|?<8ySs}|Wbrw05{#s1f@DtaP$q7|567dySOsG@HC8P3n?eYU=s+-mpBoE(VtNR^SYk6s$a}2}+W3_q0Dar${5Kkee=I=qKUX<@&3%1GUHI zbUbSsS!&-@RO6amr(G?~JJt&0+a<#%+tN;5KPnIeJrG2J9pHe{03n X=;GDDm{n~?3A-LEJjcTWj^*W4cyXs@ literal 0 HcmV?d00001 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 0000000000000000000000000000000000000000..d2eba88877108eb9fcd574152c9720d9124c170c GIT binary patch literal 113 zcmdml2M5v zMc>52Ak{1>#oW};BrU}#G1Vx|&@d&@ED^{zHcvLPU<4`s$;9%LiGeX8sq&7+S_y{U He}38kAK4&7 literal 0 HcmV?d00001 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 0000000000000000000000000000000000000000..0f5426c04addc6f2ec890c22c637f378738a6ba3 GIT binary patch literal 343 zcmV-d0jT~A+aTUo0RR91000000VuUE{a_RW-3b6D5jrk_H_ZVv03tD>=}gEJ^jHi3 zy*sh{`>}s(p-t9@37t4==ifcVWCBPSgoW`j$}s>d04o6R2aeDvU*-8YuZgnVvhLn= zm?S=tcFXQ6E)0qbgK11#K%%4$i#fT%Fu_29z$$hF5abPOeemm%pNs8Sr>&i=@ZqMN z6Q7eUHiKG{qECj8)AQ>+;rJ9i8XPM%zC)F9W!G$Sghp=-)n;Q5?c#fSM0b@|)H-ZR z!>6W@oAt>Sysz>^rDUI8hCfrw9KNKPRB#^PF-@(nn2ZH`fLZW|1bIW~^5tj$4d*kH4u p29tFVn=s=?8nZ!E(P9JPK-f|VyxQau(uHcEG{tj@IR@J*XDHD*nsNXD literal 0 HcmV?d00001 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 0000000000000000000000000000000000000000..dc3b09d1bcba56661f9ea0874df1c663aa40b485 GIT binary patch literal 377 zcmV-<0fznz+aTU~0RR91000000VuUE{a{!F?Fs;P5f*CUjN=wrv9X zWZx6HZCmU{h}PzHj#ipA|?<8ySs}|Wbrw05{#s1f@DtaP$q7|567dySOsG@HC8P3n?eYU=s+-mpBoE(VtNR^SYk6s$a}2}+W3_q0Dar${5Kkee=I=qKUX<@&3%1GUHI zbUbSsS!&-@RO6amr(G?~JJt&0+a<#%+tN;5KPnIeJrG2J9pHe{03n X=;GDDm{n~?3A-LEJjcTWj^*W4cyXs@ literal 0 HcmV?d00001 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 0000000000000000000000000000000000000000..d610a0fd3fff9aedfb16b97d8c1e10bae28d26e3 GIT binary patch literal 293 zcmV+=0owiy+aTU00RR91000000VuUE{b1w*O$Y#D4=^X7CCvf<|NsA>p=N*D?y^3S zw)>OTTF55b_P>xNTK8KF>im1fiWVs)sZ2v;9UZ;VZruj?Xh41E)5UtvAxA$X$?lC@ zDz)uuRpN7YNUYg9*v9TC@w~~%Na}AF`;Im~Mbw!Y7q;kmT-hkTPfGW$&gvR3LOrr= zP`;_iHb*^S*C69v%`WJ9>zm2;)t6B|dJ+T(^)*l>GIUaVCJtf9?x;7{$y#_S?9Oo% z8ifWkC>X{9L$@*nGyvcr97h@ez>^rDUV)Su->_UW1HTrBuFWw@&TWI?F~`O+mbE!( r4ZCa>)L^m>ViRWkNMkmLDq3tH90+46sU9IysJ2E^5U|ZVBWH>Otrdio literal 0 HcmV?d00001 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 0000000000000000000000000000000000000000..3770a746a55e482b2c14b6c3366eccef05074de3 GIT binary patch literal 205 zcmV;;05bm!+d3-E00000000000VuUE{UE@V1pwlTC@qjF>C4(at%ITFQP&=&X#aOC z2wQ2SauZUWU+nNQP)g2nQdcMkIHtBN|e({$CO)bv!S<8 zb97fzy!kc~{coB8#Pt7+BVwtb08nm$AfXT;Fo8i}pph`(K+zBc04t%gMk;)W7&w?` HyQp|dm+W81 literal 0 HcmV?d00001 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 0000000000000000000000000000000000000000..dc3b09d1bcba56661f9ea0874df1c663aa40b485 GIT binary patch literal 377 zcmV-<0fznz+aTU~0RR91000000VuUE{a{!F?Fs;P5f*CUjN=wrv9X zWZx6HZCmU{h}PzHj#ipA|?<8ySs}|Wbrw05{#s1f@DtaP$q7|567dySOsG@HC8P3n?eYU=s+-mpBoE(VtNR^SYk6s$a}2}+W3_q0Dar${5Kkee=I=qKUX<@&3%1GUHI zbUbSsS!&-@RO6amr(G?~JJt&0+a<#%+tN;5KPnIeJrG2J9pHe{03n X=;GDDm{n~?3A-LEJjcTWj^*W4cyXs@ literal 0 HcmV?d00001 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 0000000000000000000000000000000000000000..34a0321afeeaa5ad257b714c1805626e8b32a8c3 GIT binary patch literal 113 zcmdml2M5v zMc*RDGSSF1DcQo<($df($s{c~&C)12&D1E-+|bA**^&{Y^d}R`PbLQ6i&;{4E!Ijf J{Lylb1pq0WAiDqn literal 0 HcmV?d00001 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 0000000000000000000000000000000000000000..d610a0fd3fff9aedfb16b97d8c1e10bae28d26e3 GIT binary patch literal 293 zcmV+=0owiy+aTU00RR91000000VuUE{b1w*O$Y#D4=^X7CCvf<|NsA>p=N*D?y^3S zw)>OTTF55b_P>xNTK8KF>im1fiWVs)sZ2v;9UZ;VZruj?Xh41E)5UtvAxA$X$?lC@ zDz)uuRpN7YNUYg9*v9TC@w~~%Na}AF`;Im~Mbw!Y7q;kmT-hkTPfGW$&gvR3LOrr= zP`;_iHb*^S*C69v%`WJ9>zm2;)t6B|dJ+T(^)*l>GIUaVCJtf9?x;7{$y#_S?9Oo% z8ifWkC>X{9L$@*nGyvcr97h@ez>^rDUV)Su->_UW1HTrBuFWw@&TWI?F~`O+mbE!( r4ZCa>)L^m>ViRWkNMkmLDq3tH90+46sU9IysJ2E^5U|ZVBWH>Otrdio literal 0 HcmV?d00001 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 0000000000000000000000000000000000000000..0f5426c04addc6f2ec890c22c637f378738a6ba3 GIT binary patch literal 343 zcmV-d0jT~A+aTUo0RR91000000VuUE{a_RW-3b6D5jrk_H_ZVv03tD>=}gEJ^jHi3 zy*sh{`>}s(p-t9@37t4==ifcVWCBPSgoW`j$}s>d04o6R2aeDvU*-8YuZgnVvhLn= zm?S=tcFXQ6E)0qbgK11#K%%4$i#fT%Fu_29z$$hF5abPOeemm%pNs8Sr>&i=@ZqMN z6Q7eUHiKG{qECj8)AQ>+;rJ9i8XPM%zC)F9W!G$Sghp=-)n;Q5?c#fSM0b@|)H-ZR z!>6W@oAt>Sysz>^rDUI8hCfrw9KNKPRB#^PF-@(nn2ZH`fLZW|1bIW~^5tj$4d*kH4u p29tFVn=s=?8nZ!E(P9JPK-f|VyxQau(uHcEG{tj@IR@J*XDHD*nsNXD literal 0 HcmV?d00001 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 0000000000000000000000000000000000000000..dc3b09d1bcba56661f9ea0874df1c663aa40b485 GIT binary patch literal 377 zcmV-<0fznz+aTU~0RR91000000VuUE{a{!F?Fs;P5f*CUjN=wrv9X zWZx6HZCmU{h}PzHj#ipA|?<8ySs}|Wbrw05{#s1f@DtaP$q7|567dySOsG@HC8P3n?eYU=s+-mpBoE(VtNR^SYk6s$a}2}+W3_q0Dar${5Kkee=I=qKUX<@&3%1GUHI zbUbSsS!&-@RO6amr(G?~JJt&0+a<#%+tN;5KPnIeJrG2J9pHe{03n X=;GDDm{n~?3A-LEJjcTWj^*W4cyXs@ literal 0 HcmV?d00001 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 0000000000000000000000000000000000000000..2e49f05dd5c76e30903126df9334e98a63511388 GIT binary patch literal 204 zcmV;-05ks#+d3-D00000000000VuUE{UE@N1ps1;DJ+mF>BpNuW)q5%iXZX9)@H|n zQuoP3?hTS}9h~2Cw=FkmFj*Ta6y%E#Nm6&Dot((U#F2Iam<%G0v=kEs2oL`o2C;$^ zO!$urgA22?JrI6aMav@I)b)66smMhMCFg4p^*ENj_R(zBG`V&6vgKHloiel@Q;a=M z_Zn)|t#4`Psy2N%-v34;0Fn9srEyRkBn-(J1}YW=6bK+H2N9E7NHhcj0IX3AI|x!- Gymo?~uv5(d literal 0 HcmV?d00001 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 0000000000000000000000000000000000000000..dc3b09d1bcba56661f9ea0874df1c663aa40b485 GIT binary patch literal 377 zcmV-<0fznz+aTU~0RR91000000VuUE{a{!F?Fs;P5f*CUjN=wrv9X zWZx6HZCmU{h}PzHj#ipA|?<8ySs}|Wbrw05{#s1f@DtaP$q7|567dySOsG@HC8P3n?eYU=s+-mpBoE(VtNR^SYk6s$a}2}+W3_q0Dar${5Kkee=I=qKUX<@&3%1GUHI zbUbSsS!&-@RO6amr(G?~JJt&0+a<#%+tN;5KPnIeJrG2J9pHe{03n X=;GDDm{n~?3A-LEJjcTWj^*W4cyXs@ literal 0 HcmV?d00001 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 0000000000000000000000000000000000000000..0b6ad009b82b602aae8c5a6363bea60bf240b085 GIT binary patch literal 113 zcmdp=N*D?y^3S zw)>OTTF55b_P>xNTK8KF>im1fiWVs)sZ2v;9UZ;VZruj?Xh41E)5UtvAxA$X$?lC@ zDz)uuRpN7YNUYg9*v9TC@w~~%Na}AF`;Im~Mbw!Y7q;kmT-hkTPfGW$&gvR3LOrr= zP`;_iHb*^S*C69v%`WJ9>zm2;)t6B|dJ+T(^)*l>GIUaVCJtf9?x;7{$y#_S?9Oo% z8ifWkC>X{9L$@*nGyvcr97h@ez>^rDUV)Su->_UW1HTrBuFWw@&TWI?F~`O+mbE!( r4ZCa>)L^m>ViRWkNMkmLDq3tH90+46sU9IysJ2E^5U|ZVBWH>Otrdio literal 0 HcmV?d00001 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 0000000000000000000000000000000000000000..dc3b09d1bcba56661f9ea0874df1c663aa40b485 GIT binary patch literal 377 zcmV-<0fznz+aTU~0RR91000000VuUE{a{!F?Fs;P5f*CUjN=wrv9X zWZx6HZCmU{h}PzHj#ipA|?<8ySs}|Wbrw05{#s1f@DtaP$q7|567dySOsG@HC8P3n?eYU=s+-mpBoE(VtNR^SYk6s$a}2}+W3_q0Dar${5Kkee=I=qKUX<@&3%1GUHI zbUbSsS!&-@RO6amr(G?~JJt&0+a<#%+tN;5KPnIeJrG2J9pHe{03n X=;GDDm{n~?3A-LEJjcTWj^*W4cyXs@ literal 0 HcmV?d00001 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 0000000000000000000000000000000000000000..6dd2ee6437ca27d2cca3169bfe1686fb67270b36 GIT binary patch literal 166 zcmd?=#FGzIDPF2`_xB~R>uxJos#HjRL0Q2(QxRxJ_{pAB@>g9 z0s}*ezF|^Ia%z%cqN!=Bagw=#shOpzk-4RrQJRUFX;N~Mv8kc4WonX{v4wF`qJ?3K zxw(;{nQ^kIg^7WosiC1+8WSTE&@`ZtOPPK$F|mM128Nm?wJ+~ltd-anrv48^Gt5~n G(EtELEib|V literal 0 HcmV?d00001 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 0000000000000000000000000000000000000000..dc3b09d1bcba56661f9ea0874df1c663aa40b485 GIT binary patch literal 377 zcmV-<0fznz+aTU~0RR91000000VuUE{a{!F?Fs;P5f*CUjN=wrv9X zWZx6HZCmU{h}PzHj#ipA|?<8ySs}|Wbrw05{#s1f@DtaP$q7|567dySOsG@HC8P3n?eYU=s+-mpBoE(VtNR^SYk6s$a}2}+W3_q0Dar${5Kkee=I=qKUX<@&3%1GUHI zbUbSsS!&-@RO6amr(G?~JJt&0+a<#%+tN;5KPnIeJrG2J9pHe{03n X=;GDDm{n~?3A-LEJjcTWj^*W4cyXs@ literal 0 HcmV?d00001 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 0000000000000000000000000000000000000000..910d501ea5d8f368eabcf879f7ac7a5b60d39e39 GIT binary patch literal 113 zcmd literal 0 HcmV?d00001 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 0000000000000000000000000000000000000000..0f5426c04addc6f2ec890c22c637f378738a6ba3 GIT binary patch literal 343 zcmV-d0jT~A+aTUo0RR91000000VuUE{a_RW-3b6D5jrk_H_ZVv03tD>=}gEJ^jHi3 zy*sh{`>}s(p-t9@37t4==ifcVWCBPSgoW`j$}s>d04o6R2aeDvU*-8YuZgnVvhLn= zm?S=tcFXQ6E)0qbgK11#K%%4$i#fT%Fu_29z$$hF5abPOeemm%pNs8Sr>&i=@ZqMN z6Q7eUHiKG{qECj8)AQ>+;rJ9i8XPM%zC)F9W!G$Sghp=-)n;Q5?c#fSM0b@|)H-ZR z!>6W@oAt>Sysz>^rDUI8hCfrw9KNKPRB#^PF-@(nn2ZH`fLZW|1bIW~^5tj$4d*kH4u p29tFVn=s=?8nZ!E(P9JPK-f|VyxQau(uHcEG{tj@IR@J*XDHD*nsNXD literal 0 HcmV?d00001 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 0000000000000000000000000000000000000000..d610a0fd3fff9aedfb16b97d8c1e10bae28d26e3 GIT binary patch literal 293 zcmV+=0owiy+aTU00RR91000000VuUE{b1w*O$Y#D4=^X7CCvf<|NsA>p=N*D?y^3S zw)>OTTF55b_P>xNTK8KF>im1fiWVs)sZ2v;9UZ;VZruj?Xh41E)5UtvAxA$X$?lC@ zDz)uuRpN7YNUYg9*v9TC@w~~%Na}AF`;Im~Mbw!Y7q;kmT-hkTPfGW$&gvR3LOrr= zP`;_iHb*^S*C69v%`WJ9>zm2;)t6B|dJ+T(^)*l>GIUaVCJtf9?x;7{$y#_S?9Oo% z8ifWkC>X{9L$@*nGyvcr97h@ez>^rDUV)Su->_UW1HTrBuFWw@&TWI?F~`O+mbE!( r4ZCa>)L^m>ViRWkNMkmLDq3tH90+46sU9IysJ2E^5U|ZVBWH>Otrdio literal 0 HcmV?d00001 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 0000000000000000000000000000000000000000..dc3b09d1bcba56661f9ea0874df1c663aa40b485 GIT binary patch literal 377 zcmV-<0fznz+aTU~0RR91000000VuUE{a{!F?Fs;P5f*CUjN=wrv9X zWZx6HZCmU{h}PzHj#ipA|?<8ySs}|Wbrw05{#s1f@DtaP$q7|567dySOsG@HC8P3n?eYU=s+-mpBoE(VtNR^SYk6s$a}2}+W3_q0Dar${5Kkee=I=qKUX<@&3%1GUHI zbUbSsS!&-@RO6amr(G?~JJt&0+a<#%+tN;5KPnIeJrG2J9pHe{03n X=;GDDm{n~?3A-LEJjcTWj^*W4cyXs@ literal 0 HcmV?d00001 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 0000000000000000000000000000000000000000..3f96c21d1d2d743aa947ebce9674d9b3a601b66c GIT binary patch literal 199 zcmV;&066~)+d3-800000000000VuUE{UE+&1pq=zC@YXD=?{pAAdVZNcg3V58kKJI zEU4|aUtyvEj#6&9<=$JoN1;gMtKRk|DQ`49amawyU~!SqXn4ZVP?-M=bWY_+4#fZS z 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()