Method save_to_dir

axlearn/common/checkpointer.py:314–328

Starts a save to the given directories. The writes may happen in the background and not finish until wait_until_finished(). Storage implementations should invoke on_commit_callback(ckpt_dir, commit_data) when async write has completed.

(
        self,
        *,
        step: int,
        state: NestedTensor,
        ckpt_dir: str,
        on_commit_callback: StateStorageCommitCallback,
    )

    """Base StateStorage."""

    def save_to_dir(
        self,
        *,
        step: int,
        state: NestedTensor,
        ckpt_dir: str,
        on_commit_callback: StateStorageCommitCallback,
    ):
        """Starts a save to the given directories.

        The writes may happen in the background and not finish until wait_until_finished().
        Storage implementations should invoke on_commit_callback(ckpt_dir, commit_data) when async
        write has completed.
        """
        raise NotImplementedError(type(self))

    def wait_until_finished(self):
        """Waits for async writes to finish."""