feat: add checkpoint migration registry

CHANGELOG.md

@@ -9,6 +9,15 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 
 ### Added
 
+- Add a checkpoint migration registry (`neural_lam/migrations.py`) so
+  checkpoints from older neural-lam versions keep loading across renames
+  and restructuring of the model classes, instead of failing outright.
+  Checkpoints are now stamped with an internal schema version;
+  `ForecasterModule.on_load_checkpoint` migrates old ones automatically
+  with a warning, and a standalone `python -m neural_lam.migrate_checkpoint`
+  script can upgrade a checkpoint file in place.
+  [\#48](https://github.com/mllam/neural-lam/issues/48) @Sharkyii
+
 - Add `PropagationNet` GNN layer that incentivises directional message
   propagation from sender to receiver nodes, and expose it alongside
   `InteractionNet` through four new CLI arguments (`--g2m_gnn_type`,

neural_lam/migrate_checkpoint.py

@@ -0,0 +1,93 @@
+"""Standalone script to upgrade a checkpoint's state dict in place.
+
+Wraps the same `neural_lam.migrations.apply_migrations` used by
+`ForecasterModule.on_load_checkpoint`, so a checkpoint can be migrated
+once and re-saved instead of being migrated in memory on every load. See
+issue #48 for the motivation.
+"""
+
+# Standard library
+import os
+from argparse import ArgumentDefaultsHelpFormatter, ArgumentParser
+from typing import Optional
+
+# Third-party
+import torch
+
+# Local
+from . import migrations
+
+
+def migrate_checkpoint_file(load_path: str, save_path: str) -> None:
+    """
+    Migrate the checkpoint at `load_path` to the current state dict schema
+    version and save the result to `save_path`.
+
+    Parameters
+    ----------
+    load_path : str
+        Path to the checkpoint file to migrate.
+    save_path : str
+        Path to save the migrated checkpoint to.
+    """
+    checkpoint = torch.load(load_path, map_location="cpu", weights_only=False)
+    checkpoint_version = checkpoint.get("neural_lam_checkpoint_version", 0)
+
+    if checkpoint_version >= migrations.CURRENT_CHECKPOINT_VERSION:
+        print(
+            f"Checkpoint is already at version {checkpoint_version}, "
+            "nothing to migrate."
+        )
+        return
+
+    checkpoint["state_dict"], checkpoint_version = migrations.apply_migrations(
+        checkpoint["state_dict"], checkpoint_version
+    )
+    checkpoint["neural_lam_checkpoint_version"] = checkpoint_version
+
+    torch.save(checkpoint, save_path)
+    print(
+        f"Migrated checkpoint to version {checkpoint_version}, "
+        f"saved to {save_path}"
+    )
+
+
+def cli(input_args: Optional[list[str]] = None) -> None:
+    """
+    Parse CLI arguments and call `migrate_checkpoint_file`.
+
+    Parameters
+    ----------
+    input_args : list[str] or None, optional
+        Argument list forwarded to :class:`argparse.ArgumentParser`. When
+        ``None``, ``sys.argv`` is used.
+    """
+    parser = ArgumentParser(
+        description="Upgrade a neural-lam checkpoint's state dict",
+        formatter_class=ArgumentDefaultsHelpFormatter,
+    )
+    parser.add_argument(
+        "--load",
+        type=str,
+        required=True,
+        help="Path to checkpoint file to migrate",
+    )
+    parser.add_argument(
+        "--save",
+        type=str,
+        default=None,
+        help="Path to save migrated checkpoint to. Defaults to "
+        "'upgraded_<load file name>' next to the input file.",
+    )
+    args = parser.parse_args(input_args)
+
+    save_path = args.save
+    if save_path is None:
+        load_dirname, load_basename = os.path.split(args.load)
+        save_path = os.path.join(load_dirname, f"upgraded_{load_basename}")
+
+    migrate_checkpoint_file(load_path=args.load, save_path=save_path)
+
+
+if __name__ == "__main__":
+    cli()

neural_lam/migrations.py

@@ -0,0 +1,158 @@
+"""Registry of checkpoint state-dict migrations.
+
+Handles the case where a checkpoint was saved by an older version of
+neural-lam whose model classes have since been renamed, restructured or
+removed. See issue #48 for the background discussion.
+
+Checkpoints are stamped with an internal, monotonically increasing
+``neural_lam_checkpoint_version`` (unrelated to the package's own semver,
+since that is derived from git tags and is not meaningful for dev
+installs). Each entry in ``MIGRATIONS`` describes how to bring a state
+dict from the version before it up to its ``target_version``.
+"""
+
+# Standard library
+from dataclasses import dataclass
+from typing import Callable, Optional
+
+StateDict = dict
+
+
+@dataclass(frozen=True)
+class Migration:
+    """A single step that brings a state dict up to ``target_version``."""
+
+    target_version: int
+    description: str
+    apply: Optional[Callable[[StateDict], StateDict]] = None
+    hard_break: bool = False
+
+    def __post_init__(self) -> None:
+        """Validate that exactly one of `apply` or `hard_break` is set."""
+        if self.hard_break == (self.apply is not None):
+            raise ValueError(
+                "Migration must set exactly one of `apply` or `hard_break`"
+            )
+
+
+def rename_keys(
+    old_prefix: str, new_prefix: str
+) -> Callable[[StateDict], StateDict]:
+    """
+    Build a migration function that renames all state dict keys starting
+    with ``old_prefix`` to start with ``new_prefix`` instead.
+    """
+
+    def _apply(state_dict: StateDict) -> StateDict:
+        """Rename keys starting with `old_prefix` to start with `new_prefix`."""
+        for old_key in [
+            key for key in state_dict if key.startswith(old_prefix)
+        ]:
+            new_key = new_prefix + old_key[len(old_prefix) :]
+            state_dict[new_key] = state_dict.pop(old_key)
+        return state_dict
+
+    return _apply
+
+
+def remove_keys(keys: list[str]) -> Callable[[StateDict], StateDict]:
+    """Build a migration function that drops the given keys, if present."""
+
+    def _apply(state_dict: StateDict) -> StateDict:
+        """Drop `keys` from `state_dict`, if present."""
+        for key in keys:
+            state_dict.pop(key, None)
+        return state_dict
+
+    return _apply
+
+
+def _prefix_flat_ar_model_keys(state_dict: StateDict) -> StateDict:
+    """
+    Move every parameter of the pre-refactor, flat ``ARModel`` under
+    ``forecaster.predictor.``, matching the ``ForecasterModule`` /
+    ``Forecaster`` / ``StepPredictor`` split introduced in #208.
+    """
+    unprefixed_keys = ("interior_mask_bool", "per_var_std")
+    for old_key in [
+        key
+        for key in state_dict
+        if not key.startswith("forecaster.") and key not in unprefixed_keys
+    ]:
+        state_dict[f"forecaster.predictor.{old_key}"] = state_dict.pop(old_key)
+    return state_dict
+
+
+# Ordered oldest to newest. Each migration is applied to a state dict that
+# is already at `target_version - 1`ish (in practice: at any version below
+# `target_version`), so migrations must remain valid to run in sequence.
+MIGRATIONS: list[Migration] = [
+    Migration(
+        target_version=1,
+        description=(
+            "Rename 'g2m_gnn.grid_mlp' to 'encoding_grid_mlp', from moving "
+            "the grid MLP out of the InteractionNet class."
+        ),
+        apply=rename_keys("g2m_gnn.grid_mlp", "encoding_grid_mlp"),
+    ),
+    Migration(
+        target_version=2,
+        description=(
+            "Prefix flat ARModel parameters with 'forecaster.predictor.', "
+            "from splitting ARModel into ForecasterModule/Forecaster/"
+            "StepPredictor (#208)."
+        ),
+        apply=_prefix_flat_ar_model_keys,
+    ),
+]
+
+CURRENT_CHECKPOINT_VERSION = max(
+    (migration.target_version for migration in MIGRATIONS), default=0
+)
+
+
+def apply_migrations(
+    state_dict: StateDict, from_version: int
+) -> tuple[StateDict, int]:
+    """
+    Migrate ``state_dict`` from ``from_version`` up to
+    ``CURRENT_CHECKPOINT_VERSION``, applying each intervening migration in
+    order.
+
+    Parameters
+    ----------
+    state_dict : dict
+        The checkpoint's state dict, migrated in place and also returned.
+    from_version : int
+        The checkpoint's ``neural_lam_checkpoint_version`` (0 if the
+        checkpoint predates this versioning scheme entirely).
+
+    Returns
+    -------
+    tuple[dict, int]
+        The migrated state dict, and the version it now conforms to. Equal
+        to ``CURRENT_CHECKPOINT_VERSION`` unless a hard break stopped the
+        chain early (which raises rather than returning).
+
+    Raises
+    ------
+    RuntimeError
+        If a migration in the chain is a hard break, i.e. one that cannot
+        convert old parameters into new ones automatically.
+    """
+    reached_version = from_version
+    for migration in sorted(MIGRATIONS, key=lambda m: m.target_version):
+        if migration.target_version <= reached_version:
+            continue
+        if migration.hard_break:
+            raise RuntimeError(
+                f"Checkpoint at version {reached_version} cannot be "
+                f"automatically migrated past version "
+                f"{migration.target_version}: {migration.description}. "
+                "Please re-train, or convert it using the neural-lam "
+                "version just before this change."
+            )
+        assert migration.apply is not None  # guaranteed by __post_init__
+        state_dict = migration.apply(state_dict)
+        reached_version = migration.target_version
+    return state_dict, reached_version

neural_lam/models/module.py

@@ -17,7 +17,7 @@
 from neural_lam.utils import get_integer_time
 
 # Local
-from .. import metrics, vis
+from .. import metrics, migrations, vis
 from ..config import NeuralLAMConfig
 from ..datastore import BaseDatastore
 from ..loss_weighting import get_state_feature_weighting
@@ -961,6 +961,22 @@ def on_test_epoch_end(self):
         # otherwise stay permanently False).
         self.plotted_examples = 0
 
+    def on_save_checkpoint(self, checkpoint):
+        """
+        Stamp the checkpoint with the state dict schema version it was
+        saved with, so `on_load_checkpoint` knows which migrations (if
+        any) it needs to apply when this checkpoint is loaded again in a
+        future version of neural-lam.
+
+        Parameters
+        ----------
+        checkpoint : dict
+            The checkpoint dictionary about to be saved.
+        """
+        checkpoint["neural_lam_checkpoint_version"] = (
+            migrations.CURRENT_CHECKPOINT_VERSION
+        )
+
     def on_load_checkpoint(self, checkpoint):
         """
         Perform actions when loading a checkpoint.
@@ -971,43 +987,20 @@ def on_load_checkpoint(self, checkpoint):
         checkpoint : dict
             The loaded checkpoint dictionary.
         """
-        loaded_state_dict = checkpoint["state_dict"]
-
-        # 1. Broad namespace remap: for pre-refactor checkpoints
-        # The old ``ARModel`` was a flat LightningModule. Everything that
-        # belonged to the predictor needs to be moved to
-        # 'forecaster.predictor.'
-        old_keys = list(loaded_state_dict.keys())
-        for key in old_keys:
-            if not key.startswith("forecaster.") and key not in (
-                "interior_mask_bool",
-                "per_var_std",
-            ):
-                new_key = f"forecaster.predictor.{key}"
-                loaded_state_dict[new_key] = loaded_state_dict.pop(key)
-
-        # 2. Specific rename from g2m_gnn.grid_mlp -> encoding_grid_mlp
-        # Will be under forecaster.predictor due to the remap above, or
-        # already there if from a recent checkpoint before this rename.
-        if (
-            "forecaster.predictor.g2m_gnn.grid_mlp.0.weight"
-            in loaded_state_dict
-        ):
-            replace_keys = list(
-                filter(
-                    lambda key: key.startswith(
-                        "forecaster.predictor.g2m_gnn.grid_mlp"
-                    ),
-                    loaded_state_dict.keys(),
-                )
+        checkpoint_version = checkpoint.get("neural_lam_checkpoint_version", 0)
+        if checkpoint_version < migrations.CURRENT_CHECKPOINT_VERSION:
+            warnings.warn(
+                f"Checkpoint was saved with state dict schema version "
+                f"{checkpoint_version}, migrating to "
+                f"{migrations.CURRENT_CHECKPOINT_VERSION}. See "
+                "neural_lam/migrations.py for details."
             )
-            for old_key in replace_keys:
-                new_key = old_key.replace(
-                    "forecaster.predictor.g2m_gnn.grid_mlp",
-                    "forecaster.predictor.encoding_grid_mlp",
+            checkpoint["state_dict"], checkpoint_version = (
+                migrations.apply_migrations(
+                    checkpoint["state_dict"], checkpoint_version
                 )
-                loaded_state_dict[new_key] = loaded_state_dict[old_key]
-                del loaded_state_dict[old_key]
+            )
+            checkpoint["neural_lam_checkpoint_version"] = checkpoint_version
 
         if not self.restore_opt:
             opt = self.configure_optimizers()