Method test_garbage_collection

axlearn/common/checkpointer_test.py:759–834

(
        self,
        ckpt_paths: Optional[Sequence[str]],
        expect_restore_step: int,
        expect_saved_steps: Sequence[int],
        listdir_add_trailing_slash: bool,
    )

            ),
        ),
        listdir_add_trailing_slash=[True, False],
    )
    def test_garbage_collection(
        self,
        ckpt_paths: Optional[Sequence[str]],
        expect_restore_step: int,
        expect_saved_steps: Sequence[int],
        listdir_add_trailing_slash: bool,
    ):
        mesh_shape = (1, 1)
        if not test_utils.is_supported_mesh_shape(mesh_shape):
            return

        orig_listdir = listdir

        def patch_tf_io_behavior(*args):
            out = orig_listdir(*args)
            return [x + "/" for x in out if not x.endswith("/")]

        # pylint: disable=line-too-long
        with (
            _mesh(mesh_shape),
            (
                mock.patch("axlearn.common.file_system.listdir", patch_tf_io_behavior)
                if listdir_add_trailing_slash
                else nullcontext()
            ),
            tempfile.TemporaryDirectory() as temp_dir,
        ):
            cfg = Checkpointer.default_config().set(
                name="test",
                dir=temp_dir,
                keep_last_n=3,
                keep_every_n_steps=2,
                gc_loop_interval_seconds=1,
            )
            cfg.save_policy.min_step = 0

            # Running gc for non-existent dir shouldn't fail.
            ckpt_fake = cfg.clone(dir=os.path.join(temp_dir, "fake_dir")).instantiate(parent=None)
            ckpt_fake._run_garbage_collection()

            ckpt: Checkpointer = cfg.instantiate(parent=None)
            state = dict(x=jnp.zeros([], dtype=jnp.int32))

            for step in range(10):
                ckpt.save(step=step, state=state)
            ckpt.wait_until_finished()

            # Mock out the checkpoints that are committed.
            if ckpt_paths:
                ckpt_paths = [os.path.join(temp_dir, f"step_{i:08d}") for i in ckpt_paths]
            else:
                ckpt_paths = Checkpointer.checkpoint_paths(cfg.dir)

            with mock.patch.object(Checkpointer, "checkpoint_paths", return_value=ckpt_paths):
                ckpt._run_garbage_collection()