(self, checkpointer_cls)
| 658 | self.assertEqual(b"123", restored_x.get_state()["values"]) |
| 659 | ckpt.stop() |
| 660 | |
| 661 | @parameterized.parameters([Checkpointer, OrbaxCheckpointer]) |
| 662 | def test_grain(self, checkpointer_cls): |
| 663 | self.skipTest("TODO(mark-b-lee): figure out why it fails on CI.") |
| 664 | if not _GRAIN_INSTALLED: |
| 665 | self.skipTest("Cannot run when grain is not installed.") |
| 666 | mesh_shape = (1, 1) |
| 667 | if not test_utils.is_supported_mesh_shape(mesh_shape): |
| 668 | return |
| 669 | with _mesh(mesh_shape): |
| 670 | cfg = _checkpointer_config(checkpointer_cls) |
| 671 | ckpt: Checkpointer = cfg.instantiate(parent=None) |
| 672 | # pylint: disable-next=possibly-used-before-assignment |
| 673 | ds = iter(range_dataset(start=1, stop=4)) |
| 674 | # Move the input_iter. |
| 675 | self.assertEqual(next(ds), 1) |
| 676 | state0 = dict(x=jnp.ones([3, 2]), y=ds) |
| 677 | |
| 678 | self.assertEqual([], os.listdir(cfg.dir)) |
| 679 | |
| 680 | ckpt.save(step=100, state=state0) |
| 681 | ckpt.wait_until_finished() |
| 682 | |
| 683 | # Check that input iterators are saved under a per-worker path. |
| 684 | # E.g., /path/to/<step>/[state/]python_0/input_iter.index. |
| 685 | state_dir = ckpt.ckpt_dir(100) |
| 686 | if "state" in os.listdir(state_dir): |
| 687 | state_dir = os.path.join(state_dir, "state") |
| 688 | self.assertIn("python_0", os.listdir(state_dir)) |
| 689 | |
| 690 | state0_specs = dict( |
| 691 | x=utils.TensorSpec(shape=[3, 2], dtype=jnp.float32), |
| 692 | # The same iterator, but with the position at 0. |
| 693 | y=iter(range_dataset(start=0, stop=4)), |
| 694 | ) |
| 695 | |
| 696 | def tensors_only(tree): |
| 697 | return ( |
| 698 | utils.prune_tree( |
| 699 | tree, should_prune=lambda _, v: not isinstance(v, utils.Tensor) |
| 700 | ), |
| 701 | ) |
| 702 | |
| 703 | step, restored_state = ckpt.restore(step=None, state=state0_specs) |
| 704 | self.assertEqual(100, step) |
| 705 | # The iterators will be different (despite pointing to the same values). |
| 706 | self.assertNestedEqual(tensors_only(state0), tensors_only(restored_state)) |
| 707 | # The restored_state contains the input_iter pointing to the next value. |
| 708 | self.assertEqual(list(range(1, 4)), list(restored_state["y"])) |
| 709 | ckpt.stop() |
| 710 |
nothing calls this directly
no test coverage detected