(self)
| 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 | |
| 711 | def test_cleanup_checkpoint(self): |
| 712 | # Mock the rmtree s.t. it does nothing. |
| 713 | with ( |
| 714 | mock.patch("axlearn.common.file_system.rmtree", side_effect=None), |
| 715 | tempfile.TemporaryDirectory() as temp_dir, |
| 716 | ): |
| 717 | # Create a few mock checkpoints. |
| 718 | ckpt_paths = [] |
| 719 | for step in [1, 2]: |
| 720 | ckpt_paths.append(os.path.join(temp_dir, f"step_{step:08d}")) |
| 721 | os.makedirs(ckpt_paths[-1]) |
| 722 | for file in ["test", "index"]: |
| 723 | with open(os.path.join(ckpt_paths[-1], file), "w", encoding="utf-8") as f: |
| 724 | f.write(str(step)) |
| 725 | self.assertEqual(Checkpointer.latest_checkpoint_path(temp_dir), ckpt_paths[-1]) |
| 726 | # Simulate a corrupted cleanup on the last ckpt. |
| 727 | Checkpointer.cleanup_checkpoint(ckpt_paths[-1], sync=False) |
| 728 | # Ensure that the last ckpt still has the "test" file. |
| 729 | with open(os.path.join(ckpt_paths[-1], "test"), encoding="utf-8") as f: |
| 730 | self.assertEqual("2", f.read()) |
| 731 | # Ensure that the last ckpt is considered invalid. |
| 732 | self.assertEqual(Checkpointer.latest_checkpoint_path(temp_dir), ckpt_paths[0]) |
| 733 |
nothing calls this directly
no test coverage detected