Tests that index files saved with orbax can be read with `read_index_file`.
(self)
| 29 | |
| 30 | class OrbaxCheckpointerTest(test_utils.TestCase): |
| 31 | def test_index(self): |
| 32 | """Tests that index files saved with orbax can be read with `read_index_file`.""" |
| 33 | mesh_shape = (1, 1) |
| 34 | if not test_utils.is_supported_mesh_shape(mesh_shape): |
| 35 | return |
| 36 | with _mesh(mesh_shape), tempfile.TemporaryDirectory() as temp_dir: |
| 37 | ckpt = ( |
| 38 | OrbaxCheckpointer.default_config() |
| 39 | .set(name="test", dir=temp_dir) |
| 40 | .instantiate(parent=None) |
| 41 | ) |
| 42 | step = 123 |
| 43 | state = dict(x=jnp.ones([3, 2])) |
| 44 | ckpt.save(step=step, state=state) |
| 45 | ckpt.wait_until_finished() |
| 46 | |
| 47 | ref_index = read_index_file(os.path.join(temp_dir, "step_00000123", "index")) |
| 48 | test_index = ckpt._manager.restore( |
| 49 | step=step, |
| 50 | # The input iterator is saved as part of `save_tf_savables`. |
| 51 | args=ocp.args.Composite( |
| 52 | index=ocp.args.JsonSave(ckpt._get_spec(step=step, state=state)) |
| 53 | ), |
| 54 | ) |
| 55 | self.assertEqual(ref_index, test_index["index"]) |
| 56 | |
| 57 | def test_checkpoint_manager_with_tracker_file(self): |
| 58 | """Test CheckpointManagerWithTrackerFile""" |
nothing calls this directly
no test coverage detected