|
| 1 | +# Copyright 2025 The Orbax Authors. |
| 2 | +# |
| 3 | +# Licensed under the Apache License, Version 2.0 (the "License"); |
| 4 | +# you may not use this file except in compliance with the License. |
| 5 | +# You may obtain a copy of the License at |
| 6 | +# |
| 7 | +# http://www.apache.org/licenses/LICENSE-2.0 |
| 8 | +# |
| 9 | +# Unless required by applicable law or agreed to in writing, software |
| 10 | +# distributed under the License is distributed on an "AS IS" BASIS, |
| 11 | +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. |
| 12 | +# See the License for the specific language governing permissions and |
| 13 | +# limitations under the License. |
| 14 | + |
| 15 | +"""Tests for ArrayCheckpointHandler.""" |
| 16 | + |
| 17 | +from absl import flags |
| 18 | +from absl.testing import parameterized |
| 19 | +from etils import epath |
| 20 | +import jax |
| 21 | +import numpy as np |
| 22 | +from orbax.checkpoint import test_utils |
| 23 | +from orbax.checkpoint._src.handlers import array_checkpoint_handler |
| 24 | +from orbax.checkpoint._src.multihost import multihost |
| 25 | +from orbax.checkpoint._src.serialization import type_handlers |
| 26 | +from orbax.checkpoint._src.testing import multiprocess_test |
| 27 | + |
| 28 | + |
| 29 | +SaveArgs = type_handlers.SaveArgs |
| 30 | +ArraySaveArgs = array_checkpoint_handler.ArraySaveArgs |
| 31 | +ArrayRestoreArgs = array_checkpoint_handler.ArrayRestoreArgs |
| 32 | + |
| 33 | + |
| 34 | +FLAGS = flags.FLAGS |
| 35 | + |
| 36 | + |
| 37 | +class ArrayCheckpointHandler(array_checkpoint_handler.ArrayCheckpointHandler): |
| 38 | + |
| 39 | + def save(self, directory, *args, **kwargs): |
| 40 | + super().save(directory, *args, **kwargs) |
| 41 | + test_utils.sync_global_processes('ArrayCheckpointHandler:save') |
| 42 | + if multihost.process_index() == 0: |
| 43 | + self.finalize(directory) |
| 44 | + test_utils.sync_global_processes('ArrayCheckpointHandler:finalize') |
| 45 | + |
| 46 | + |
| 47 | +class ArrayCheckpointHandlerTest( |
| 48 | + parameterized.TestCase, multiprocess_test.MultiProcessTest |
| 49 | +): |
| 50 | + |
| 51 | + def setUp(self): |
| 52 | + super().setUp() |
| 53 | + self.devices = np.asarray(jax.devices()) |
| 54 | + self.directory = epath.Path( |
| 55 | + self.create_tempdir(name='checkpointing_test').full_path |
| 56 | + ) |
| 57 | + |
| 58 | + test_utils.sync_global_processes( |
| 59 | + 'ArrayCheckpointHandlerTest:setup_complete' |
| 60 | + ) |
| 61 | + |
| 62 | + def tearDown(self): |
| 63 | + test_utils.sync_global_processes( |
| 64 | + 'ArrayCheckpointHandlerTest:tests_complete' |
| 65 | + ) |
| 66 | + super().tearDown() |
| 67 | + |
| 68 | + def validate_save(self): |
| 69 | + path = self.directory / array_checkpoint_handler.PYTREE_METADATA_FILE |
| 70 | + self.assertTrue(path.exists()) |
| 71 | + |
| 72 | + def test_array(self): |
| 73 | + checkpoint_handler = ArrayCheckpointHandler() |
| 74 | + mesh = jax.sharding.Mesh(self.devices, ('x',)) |
| 75 | + mesh_axes = jax.sharding.PartitionSpec( |
| 76 | + 'x', |
| 77 | + ) |
| 78 | + arr = test_utils.create_sharded_array(np.arange(16), mesh, mesh_axes) |
| 79 | + save_args = SaveArgs() |
| 80 | + checkpoint_handler.save(self.directory, args=ArraySaveArgs(arr, save_args)) |
| 81 | + self.validate_save() |
| 82 | + restored = checkpoint_handler.restore( |
| 83 | + self.directory, |
| 84 | + args=ArrayRestoreArgs( |
| 85 | + restore_args=type_handlers.ArrayRestoreArgs( |
| 86 | + restore_type=jax.Array, mesh=mesh, mesh_axes=mesh_axes |
| 87 | + ) |
| 88 | + ), |
| 89 | + ) |
| 90 | + test_utils.assert_tree_equal(self, [arr], [restored]) |
| 91 | + checkpoint_handler.close() |
| 92 | + |
| 93 | + def test_numpy_array(self): |
| 94 | + checkpoint_handler = ArrayCheckpointHandler() |
| 95 | + arr = np.arange(16) |
| 96 | + save_args = SaveArgs() |
| 97 | + checkpoint_handler.save(self.directory, args=ArraySaveArgs(arr, save_args)) |
| 98 | + self.validate_save() |
| 99 | + restored = checkpoint_handler.restore( |
| 100 | + self.directory, |
| 101 | + args=ArrayRestoreArgs( |
| 102 | + restore_args=type_handlers.RestoreArgs(restore_type=np.ndarray) |
| 103 | + ), |
| 104 | + ) |
| 105 | + test_utils.assert_tree_equal(self, [arr], [restored]) |
| 106 | + checkpoint_handler.close() |
| 107 | + |
| 108 | + def test_scalar(self): |
| 109 | + checkpoint_handler = ArrayCheckpointHandler() |
| 110 | + save_args = SaveArgs() |
| 111 | + checkpoint_handler.save(self.directory, args=ArraySaveArgs(5, save_args)) |
| 112 | + self.validate_save() |
| 113 | + restored = checkpoint_handler.restore( |
| 114 | + self.directory, |
| 115 | + args=ArrayRestoreArgs( |
| 116 | + restore_args=type_handlers.RestoreArgs(restore_type=int) |
| 117 | + ), |
| 118 | + ) |
| 119 | + self.assertEqual(5, restored) |
| 120 | + checkpoint_handler.close() |
| 121 | + |
| 122 | + def test_invalid_type(self): |
| 123 | + checkpoint_handler = ArrayCheckpointHandler() |
| 124 | + with self.assertRaises(TypeError): |
| 125 | + checkpoint_handler.save(self.directory, args=ArraySaveArgs('hi')) |
| 126 | + checkpoint_handler.close() |
| 127 | + |
| 128 | + def test_different_name(self): |
| 129 | + checkpoint_name = 'my_array' |
| 130 | + checkpoint_handler = ArrayCheckpointHandler(checkpoint_name=checkpoint_name) |
| 131 | + arr = np.arange(16) |
| 132 | + save_args = SaveArgs() |
| 133 | + checkpoint_handler.save(self.directory, args=ArraySaveArgs(arr, save_args)) |
| 134 | + self.validate_save() |
| 135 | + restored = checkpoint_handler.restore( |
| 136 | + self.directory, |
| 137 | + args=ArrayRestoreArgs( |
| 138 | + restore_args=type_handlers.RestoreArgs(restore_type=np.ndarray) |
| 139 | + ), |
| 140 | + ) |
| 141 | + test_utils.assert_tree_equal(self, [arr], [restored]) |
| 142 | + checkpoint_handler.close() |
| 143 | + |
| 144 | + def test_restore_type(self): |
| 145 | + pytree = 5 |
| 146 | + checkpoint_handler = ArrayCheckpointHandler() |
| 147 | + |
| 148 | + checkpoint_handler.save(self.directory, args=ArraySaveArgs(pytree)) |
| 149 | + restored = checkpoint_handler.restore( |
| 150 | + self.directory, |
| 151 | + args=ArrayRestoreArgs( |
| 152 | + restore_args=type_handlers.RestoreArgs(restore_type=np.ndarray) |
| 153 | + ), |
| 154 | + ) |
| 155 | + self.assertIsInstance(restored, np.ndarray) |
| 156 | + |
| 157 | + |
| 158 | +if __name__ == '__main__': |
| 159 | + multiprocess_test.main() |
0 commit comments