diff --git a/CHANGELOG.md b/CHANGELOG.md index fd123dee..d5043fc6 100644 --- a/CHANGELOG.md +++ b/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`, diff --git a/neural_lam/migrate_checkpoint.py b/neural_lam/migrate_checkpoint.py new file mode 100644 index 00000000..be963dfa --- /dev/null +++ b/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_' 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() diff --git a/neural_lam/migrations.py b/neural_lam/migrations.py new file mode 100644 index 00000000..28d64bf2 --- /dev/null +++ b/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 diff --git a/neural_lam/models/module.py b/neural_lam/models/module.py index 71ce7951..e2d679e4 100644 --- a/neural_lam/models/module.py +++ b/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() diff --git a/tests/test_checkpoint.py b/tests/test_checkpoint.py index 2e5f3148..667add33 100644 --- a/tests/test_checkpoint.py +++ b/tests/test_checkpoint.py @@ -7,21 +7,14 @@ # First-party from neural_lam import config as nlconfig +from neural_lam import migrations from neural_lam.create_graph import create_graph_from_datastore from neural_lam.models import ARForecaster, ForecasterModule, GraphLAM from tests.dummy_datastore import DummyDatastore -def test_saved_checkpoint_excludes_datastore_and_forecaster(tmp_path): - """ - Regression check for issue #148: heavy non-pickle-safe objects - (`datastore`, `forecaster`) must be excluded from the saved Lightning - hyperparameters so that checkpoints stay small and portable, and so - that `load_from_checkpoint` requires them to be passed in explicitly. - """ - datastore = DummyDatastore() - - # Build the minimum graph the GraphLAM predictor needs. +def _build_forecaster_module(datastore: DummyDatastore) -> ForecasterModule: + """Build the minimal `ForecasterModule` the checkpoint tests need.""" graph_dir_path = Path(datastore.root_path) / "graph" / "1level" if not graph_dir_path.exists(): create_graph_from_datastore( @@ -33,7 +26,7 @@ def test_saved_checkpoint_excludes_datastore_and_forecaster(tmp_path): config = nlconfig.NeuralLAMConfig( datastore=nlconfig.DatastoreSelection( kind=datastore.SHORT_NAME, - config_path=datastore.root_path, + config_path=str(datastore.root_path), ), ) @@ -51,7 +44,7 @@ def test_saved_checkpoint_excludes_datastore_and_forecaster(tmp_path): output_clamping_upper=config.training.output_clamping.upper, ) forecaster = ARForecaster(predictor, datastore) - model = ForecasterModule( + return ForecasterModule( forecaster=forecaster, config=config, datastore=datastore, @@ -61,6 +54,16 @@ def test_saved_checkpoint_excludes_datastore_and_forecaster(tmp_path): val_steps_to_log=[1], ) + +def test_saved_checkpoint_excludes_datastore_and_forecaster(tmp_path): + """ + Regression check for issue #148: heavy non-pickle-safe objects + (`datastore`, `forecaster`) must be excluded from the saved Lightning + hyperparameters so that checkpoints stay small and portable, and so + that `load_from_checkpoint` requires them to be passed in explicitly. + """ + model = _build_forecaster_module(DummyDatastore()) + # Lightning's in-memory hparams must already drop these. assert "datastore" not in model.hparams assert "forecaster" not in model.hparams @@ -83,3 +86,53 @@ def test_saved_checkpoint_excludes_datastore_and_forecaster(tmp_path): assert "hyper_parameters" in ckpt assert "datastore" not in ckpt["hyper_parameters"] assert "forecaster" not in ckpt["hyper_parameters"] + + +def test_saved_checkpoint_is_stamped_with_current_schema_version(tmp_path): + """`on_save_checkpoint` must stamp every new checkpoint so that a + future `on_load_checkpoint` knows whether it needs to migrate it.""" + model = _build_forecaster_module(DummyDatastore()) + + trainer = pl.Trainer( + default_root_dir=tmp_path, + accelerator="cpu", + max_epochs=0, + logger=False, + enable_checkpointing=False, + ) + trainer.strategy.connect(model) + + ckpt_path = tmp_path / "test.ckpt" + trainer.save_checkpoint(ckpt_path, weights_only=False) + + ckpt = torch.load(ckpt_path, map_location="cpu", weights_only=False) + assert ( + ckpt["neural_lam_checkpoint_version"] + == migrations.CURRENT_CHECKPOINT_VERSION + ) + + +def test_pre_refactor_checkpoint_loads_via_migration(tmp_path): + """A checkpoint from before #208 (flat `ARModel`, no version stamp, + `g2m_gnn.grid_mlp` naming) must still load, matching the current + model's state dict once migrated.""" + model = _build_forecaster_module(DummyDatastore()) + current_state_dict = model.state_dict() + + # Reverse this PR's own migrations to build a legacy state dict: strip + # the 'forecaster.predictor.' prefix and rename encoding_grid_mlp back + # to its pre-refactor name. + legacy_state_dict = {} + for key, value in current_state_dict.items(): + legacy_key = key.removeprefix("forecaster.predictor.") + legacy_key = legacy_key.replace("encoding_grid_mlp", "g2m_gnn.grid_mlp") + legacy_state_dict[legacy_key] = value + + checkpoint = {"state_dict": legacy_state_dict} + model.on_load_checkpoint(checkpoint) + + assert checkpoint["state_dict"].keys() == current_state_dict.keys() + assert ( + checkpoint["neural_lam_checkpoint_version"] + == migrations.CURRENT_CHECKPOINT_VERSION + ) diff --git a/tests/test_migrations.py b/tests/test_migrations.py new file mode 100644 index 00000000..c8876105 --- /dev/null +++ b/tests/test_migrations.py @@ -0,0 +1,120 @@ +# Third-party +import pytest + +# First-party +from neural_lam import migrations + + +def test_rename_keys_renames_matching_prefix_only(): + rename = migrations.rename_keys("old.prefix", "new.prefix") + state_dict = { + "old.prefix.weight": 1, + "old.prefix.nested.bias": 2, + "unrelated.old.prefix.weight": 3, + "other.weight": 4, + } + + result = rename(state_dict) + + assert result == { + "new.prefix.weight": 1, + "new.prefix.nested.bias": 2, + "unrelated.old.prefix.weight": 3, + "other.weight": 4, + } + + +def test_remove_keys_drops_present_keys_and_ignores_missing(): + remove = migrations.remove_keys(["a", "missing"]) + state_dict = {"a": 1, "b": 2} + + result = remove(state_dict) + + assert result == {"b": 2} + + +def test_migration_requires_exactly_one_of_apply_or_hard_break(): + with pytest.raises(ValueError): + migrations.Migration(target_version=1, description="neither") + + with pytest.raises(ValueError): + migrations.Migration( + target_version=1, + description="both", + apply=lambda sd: sd, + hard_break=True, + ) + + +def test_apply_migrations_runs_only_migrations_above_from_version(): + fake_registry = [ + migrations.Migration( + target_version=1, + description="add a", + apply=lambda sd: {**sd, "a": 1}, + ), + migrations.Migration( + target_version=2, + description="add b", + apply=lambda sd: {**sd, "b": 2}, + ), + ] + original = migrations.MIGRATIONS + migrations.MIGRATIONS = fake_registry + try: + state_dict, reached_version = migrations.apply_migrations( + {}, from_version=1 + ) + assert state_dict == {"b": 2} + assert reached_version == 2 + finally: + migrations.MIGRATIONS = original + + +def test_apply_migrations_is_a_noop_when_already_current(): + state_dict = {"x": 1} + result, reached_version = migrations.apply_migrations( + state_dict, from_version=migrations.CURRENT_CHECKPOINT_VERSION + ) + assert result == {"x": 1} + assert reached_version == migrations.CURRENT_CHECKPOINT_VERSION + + +def test_apply_migrations_raises_clear_error_on_hard_break(): + fake_registry = [ + migrations.Migration( + target_version=1, description="cannot be bridged", hard_break=True + ), + ] + original = migrations.MIGRATIONS + migrations.MIGRATIONS = fake_registry + try: + with pytest.raises(RuntimeError, match="cannot be bridged"): + migrations.apply_migrations({}, from_version=0) + finally: + migrations.MIGRATIONS = original + + +def test_pre_refactor_flat_checkpoint_migrates_to_current_layout(): + """Regression test for the two real migrations this registry replaced + in `ForecasterModule.on_load_checkpoint` (see #48).""" + legacy_state_dict = { + "g2m_gnn.grid_mlp.0.weight": "w0", + "g2m_gnn.grid_mlp.0.bias": "b0", + "some_other_layer.weight": "w1", + "interior_mask_bool": "mask", + "per_var_std": "std", + } + + state_dict, reached_version = migrations.apply_migrations( + legacy_state_dict, from_version=0 + ) + + assert reached_version == migrations.CURRENT_CHECKPOINT_VERSION + assert state_dict == { + "forecaster.predictor.encoding_grid_mlp.0.weight": "w0", + "forecaster.predictor.encoding_grid_mlp.0.bias": "b0", + "forecaster.predictor.some_other_layer.weight": "w1", + "interior_mask_bool": "mask", + "per_var_std": "std", + } diff --git a/tests/test_prediction_model_classes.py b/tests/test_prediction_model_classes.py index 73e2f905..476f736f 100644 --- a/tests/test_prediction_model_classes.py +++ b/tests/test_prediction_model_classes.py @@ -255,6 +255,8 @@ def test_forecaster_module_old_checkpoint(tmp_path): ), "config": ckpt["hyper_parameters"]["config"], } + # Remove the version stamp so the checkpoint looks pre-versioning (v0) + ckpt.pop("neural_lam_checkpoint_version", None) torch.save(ckpt, ckpt_path) # Build a fresh forecaster structure for loading weights into