From 57b427fe03eb149e41555f837bcfc3b58c71ccd7 Mon Sep 17 00:00:00 2001 From: sadamov Date: Sat, 6 Jun 2026 21:14:06 +0200 Subject: [PATCH] test: Verify datastore and forecaster are omitted from checkpoint hparams (#232) Re-applies @Jayant-kernel's PR #232 onto current main. The original test targeted the pre-#208 API (`GraphLAM(args=..., config=..., datastore=...)` all on the top-level model). Rewritten for the post-#208 layout: build a GraphLAM step predictor + ARForecaster + ForecasterModule explicitly and assert the ignore covers both heavy non-pickle-safe objects (`datastore` and the now-additional `forecaster`). Refs #148. Co-Authored-By: Jayant Co-Authored-By: Claude Opus 4.7 (1M context) --- CHANGELOG.md | 2 + tests/test_checkpoint.py | 85 ++++++++++++++++++++++++++++++++++++++++ 2 files changed, 87 insertions(+) create mode 100644 tests/test_checkpoint.py diff --git a/CHANGELOG.md b/CHANGELOG.md index fa946d02..4a7be915 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -76,6 +76,8 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 - Add unit tests for `inverse_softplus` covering roundtrip identity (parametrized over `beta`), near-zero clamping, and above-threshold linear passthrough [\#419](https://github.com/mllam/neural-lam/pull/419) @Riteesh-NITT +- Add regression test that `datastore` and `forecaster` stay excluded from the saved Lightning hyperparameters (in-memory and on-disk after a checkpoint round-trip), so `load_from_checkpoint` continues to require them to be passed in explicitly ([#148](https://github.com/mllam/neural-lam/issues/148)) [\#232](https://github.com/mllam/neural-lam/pull/232) @Jayant-kernel + - Add comprehensive type hints to `neural_lam/metrics.py` [\#447](https://github.com/mllam/neural-lam/pull/447) @sidhantpande - Add type hints to methods in `neural_lam/custom_loggers.py` [\#455](https://github.com/mllam/neural-lam/pull/455) @sidhantpande diff --git a/tests/test_checkpoint.py b/tests/test_checkpoint.py new file mode 100644 index 00000000..2e5f3148 --- /dev/null +++ b/tests/test_checkpoint.py @@ -0,0 +1,85 @@ +# Standard library +from pathlib import Path + +# Third-party +import pytorch_lightning as pl +import torch + +# First-party +from neural_lam import config as nlconfig +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. + graph_dir_path = Path(datastore.root_path) / "graph" / "1level" + if not graph_dir_path.exists(): + create_graph_from_datastore( + datastore=datastore, + output_root_path=str(graph_dir_path), + n_max_levels=1, + ) + + config = nlconfig.NeuralLAMConfig( + datastore=nlconfig.DatastoreSelection( + kind=datastore.SHORT_NAME, + config_path=datastore.root_path, + ), + ) + + predictor = GraphLAM( + datastore=datastore, + graph_name="1level", + hidden_dim=4, + hidden_layers=1, + processor_layers=1, + mesh_aggr="sum", + num_past_forcing_steps=0, + num_future_forcing_steps=0, + output_std=False, + output_clamping_lower=config.training.output_clamping.lower, + output_clamping_upper=config.training.output_clamping.upper, + ) + forecaster = ARForecaster(predictor, datastore) + model = ForecasterModule( + forecaster=forecaster, + config=config, + datastore=datastore, + loss="mse", + lr=1.0e-3, + n_example_pred=1, + val_steps_to_log=[1], + ) + + # Lightning's in-memory hparams must already drop these. + assert "datastore" not in model.hparams + assert "forecaster" not in model.hparams + + # And the on-disk checkpoint round-trip must agree. + 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) + + # In-process checkpoint, trusted source. + ckpt = torch.load(ckpt_path, map_location="cpu", weights_only=False) + assert "hyper_parameters" in ckpt + assert "datastore" not in ckpt["hyper_parameters"] + assert "forecaster" not in ckpt["hyper_parameters"]