(self)
| 600 | ckpt.stop() |
| 601 | |
| 602 | @parameterized.parameters([Checkpointer, OrbaxCheckpointer]) |
| 603 | def test_python_savable(self, checkpointer_cls): |
| 604 | mesh_shape = (1, 1) |
| 605 | if not test_utils.is_supported_mesh_shape(mesh_shape): |
| 606 | return |
| 607 | |
| 608 | class _DummySavable: |
| 609 | """A dummy class implementing PythonSavable.""" |
| 610 | |
| 611 | def __init__(self, max_step: int): |
| 612 | self._max_step = max_step |
| 613 | self._state = dict(step=0, values=b"123") |
| 614 | |
| 615 | def get_state(self): |
| 616 | return self._state |
| 617 | |
| 618 | def set_state(self, state): |
| 619 | self._state = state |
| 620 | |
| 621 | def __iter__(self): |
| 622 | for i in range(self._state["step"], self._max_step): |
| 623 | self._state["step"] += 1 |
| 624 | yield i |
| 625 | |
| 626 | with _mesh(mesh_shape): |
| 627 | cfg = _checkpointer_config(checkpointer_cls) |
| 628 | ckpt: BaseCheckpointer = cfg.instantiate(parent=None) |
| 629 | |
| 630 | x = _DummySavable(max_step=3) |
| 631 | # Check that runtime_checks is enabled. |
| 632 | self.assertIsInstance(x, PythonSavable) |
| 633 | |
| 634 | # Iterate once and save. |
| 635 | self.assertEqual(next(iter(x)), 0) |
| 636 | state0 = dict(x=x) |
| 637 | |
| 638 | self.assertEqual([], os.listdir(cfg.dir)) |
| 639 | ckpt.save(step=1, state=state0) |
| 640 | ckpt.wait_until_finished() |
| 641 | |
| 642 | # Check that input iterators are saved under a per-worker path. |
| 643 | # For Checkpointer: /path/to/<step>/python_0/x. |
| 644 | # For OrbaxCheckpointer: /path/to/<step>/python/python_0/x. |
| 645 | state_dir = ckpt.ckpt_dir(1) |
| 646 | if "python" in os.listdir(state_dir): |
| 647 | state_dir = os.path.join(state_dir, "python") |
| 648 | self.assertIn("python_0", os.listdir(state_dir)) |
| 649 | self.assertIn("x", os.listdir(os.path.join(state_dir, "python_0"))) |
| 650 | |
| 651 | # Construct a new savable object to restore. |
| 652 | state1 = dict(x=_DummySavable(max_step=3)) |
| 653 | step, restored_state = ckpt.restore(step=None, state=state1) |
| 654 | self.assertEqual(1, step) |
| 655 | restored_x: _DummySavable = restored_state["x"] |
| 656 | # The restored_state contains the iter pointing to the next value. |
| 657 | self.assertEqual(list(range(1, 3)), list(iter(restored_x))) |
| 658 | self.assertEqual(b"123", restored_x.get_state()["values"]) |
| 659 | ckpt.stop() |
nothing calls this directly
no test coverage detected