Source code for helios.chkpt_migrator

import argparse
import pathlib
import platform

import torch
import tqdm

from helios import core

from ._version import __version__
from .model.model import _InternalStateKeys
from .trainer import TrainingState, _CheckpointKeys, register_trainer_types_for_safe_load




def migrate_checkpoints_to_current_version(root: pathlib.Path) -> None:
    """
    Migrate existing checkpoints from a previous version of Helios to the current version.

    This function exists to provide backwards compatibility with checkpoints produced by
    older versions of Helios.

    Args:
        root: the root where the checkpoints are stored.
    """
    register_trainer_types_for_safe_load()

    for chkpt_path in tqdm.tqdm(
        list(root.glob("*.pth")), desc="Migrating checkpoints", unit="chkpt"
    ):
        state = core.safe_torch_load(chkpt_path)

        # 1.0 checkpoint changes
        # =======================#
        # Pre-v1.0, checkpoints didn't have a version key.
        if _CheckpointKeys.VERSION not in state:
            state[_CheckpointKeys.VERSION] = __version__

        # 1.1 checkpoint changes
        # =======================#
        # Pre-v1.1, the TrainingState struct is saved as a dictionary, not the object
        # itself.
        if isinstance(state[_CheckpointKeys.TRAINING_STATE], dict):
            state[_CheckpointKeys.TRAINING_STATE] = TrainingState(
                **state[_CheckpointKeys.TRAINING_STATE]
            )

        # 2.0 checkpoint changes
        # =======================#
        # Pre-v2.0, all model state was stored under the "model" key. v2.0 added the
        # notion of model-internal state, and user state is now saved under the "user"
        # key. Therefore everything that was under the "model" key now gets mapped to the
        # "user" sub-key.
        if _InternalStateKeys.USER not in state[_CheckpointKeys.MODEL]:
            state[_CheckpointKeys.MODEL] = {
                _InternalStateKeys.USER: state[_CheckpointKeys.MODEL]
            }

        # Pre-v2.0, the log_path and run_path were stored as separate keys. These have now
        # been moved under the "loggers" key.
        if _CheckpointKeys.LOGGERS not in state:
            loggers_state: dict[str, dict] = {}
            if "log_path" in state:
                loggers_state["root"] = {"log_file": state.pop("log_path")}
            if "run_path" in state:
                loggers_state["tensorboard"] = {"run_path": state.pop("run_path")}
            state[_CheckpointKeys.LOGGERS] = loggers_state

        # Pre-v2.0, checkpoints didn't include datamodule phase state. Default is an empty
        # dict which maps to phase 0.
        if _CheckpointKeys.DATAMODULE not in state:
            state[_CheckpointKeys.DATAMODULE] = {}

        torch.save(state, chkpt_path)



def _main() -> None:
    parser = argparse.ArgumentParser(
        description="Migration tool to convert checkpoints generated by versions of "
        f"Helios prior to {__version__}"
    )
    parser.add_argument(
        "root",
        metavar="ROOT",
        nargs=1,
        type=str,
        help="Root where the checkpoints are stored",
    )

    args = parser.parse_args()
    root = args.root[0]

    # Temporarily re-direct PosixPath to WindowsPath on Windows to avoid problems.
    tmp: type[pathlib.WindowsPath] | type[pathlib.PosixPath] | None = None
    if platform.system() == "Windows":
        tmp = pathlib.PosixPath
        pathlib.PosixPath = pathlib.WindowsPath  # type: ignore[assignment, misc]
    elif platform.system() == "Linux":
        tmp = pathlib.WindowsPath
        pathlib.WindowsPath = pathlib.PosixPath  # type: ignore[assignment, misc]

    migrate_checkpoints_to_current_version(pathlib.Path(root))

    # Restore to default values.
    if platform.system() == "Windows":
        pathlib.PosixPath = tmp  # type: ignore[assignment, misc]
    elif platform.system() == "Linux":
        pathlib.WindowsPath = tmp  # type: ignore[assignment, misc]


if __name__ == "__main__":
    _main()