Method test_save_and_restore

axlearn/common/checkpointer_test.py:85–174

(self, checkpointer_cls: Type[BaseCheckpointer])

class CheckpointerTest(test_utils.TestCase):
    @parameterized.parameters(Checkpointer, OrbaxCheckpointer)
    def test_save_and_restore(self, checkpointer_cls: Type[BaseCheckpointer]):
        mesh_shape = (1, 1)
        if not test_utils.is_supported_mesh_shape(mesh_shape):
            return
        with _mesh(mesh_shape):
            cfg = _checkpointer_config(checkpointer_cls)
            cfg.save_policy.min_step = 0
            ckpt: BaseCheckpointer = cfg.instantiate(parent=None)
            state0 = dict(x=jnp.zeros([], dtype=jnp.int32), y=jnp.ones([2], dtype=jnp.float32))
            state1 = dict(x=jnp.ones([], dtype=jnp.int32), y=jnp.ones([2], dtype=jnp.float32) + 1)

            # Restoring from an empty dir returns the input state if step=None.
            self.assertNestedEqual((None, state0), ckpt.restore(step=None, state=state0))
            self.assertNestedEqual((None, state1), ckpt.restore(step=None, state=state1))
            # With an explicit step, ValueError will be raised.
            with self.assertRaises(ValueError):
                ckpt.restore(step=0, state=state0)

            ckpt.save(step=0, state=state0)
            ckpt.wait_until_finished()
            self.assertNestedEqual((0, state0), ckpt.restore(step=0, state=state1))
            # step=None restores from the latest ckpt.
            self.assertNestedEqual((0, state0), ckpt.restore(step=None, state=state1))

            ckpt.save(step=1, state=state1)
            ckpt.wait_until_finished()
            self.assertNestedEqual((1, state1), ckpt.restore(step=1, state=state0))
            # step=None restores from the latest ckpt.
            self.assertNestedEqual((1, state1), ckpt.restore(step=None, state=state0))

            # When the given state has a different dict key: 'z' instead of 'y'.
            with self.assertRaisesRegex((ValueError, KeyError), "z"):
                ckpt.restore(
                    step=None,
                    state=dict(
                        x=jnp.zeros([], dtype=jnp.int32), z=jnp.ones([2], dtype=jnp.float32)
                    ),
                )

            # When the given state has a different array shape: [3] instead of [2] for y.
            with self.assertRaisesRegex(
                ValueError, "(checkpoint tree dtypes or shapes|not compatible)"
            ):
                ckpt.restore(
                    step=None,
                    state=dict(
                        x=jnp.zeros([], dtype=jnp.int32), y=jnp.ones([3], dtype=jnp.float32)
                    ),
                )
            # TODO(matthew_e_hopkins): revert it once upgrade jax version.
            if checkpointer_cls is Checkpointer:
                # When the given state has a different dict shape: [1] instead of [] for x.
                # Orbax throws AssertionError in this case.
                with self.assertRaisesRegex(
                    (AssertionError, ValueError),
                    "(checkpoint tree dtypes or shapes|not compatible)",
                ):
                    ckpt.restore(