Fix graph LAM GNN kwarg handling by gitcommit90 · Pull Request #688 · mllam/neural-lam

gitcommit90 · 2026-06-27T22:49:55Z

Describe your changes

Fix graph_lam startup by allowing GraphLAM to ignore hierarchical-only GNN kwargs.

Restore saved GNN type kwargs when loading checkpoints, with defaults for older checkpoints.

No new dependencies.

Issue Link

closes #686

Type of change

🐛 Bug fix (non-breaking change that fixes an issue)
✨ New feature (non-breaking change that adds functionality)
💥 Breaking change (fix or feature that would cause existing functionality to not work as expected)
📖 Documentation (Addition or improvements to documentation)

Checklist before requesting a review

My branch is up-to-date with the target branch - if not update your fork with the changes from the target branch (use pull with --rebase option if possible).
I have performed a self-review of my code
For any new/modified functions/classes I have added docstrings that clearly describe its purpose, expected inputs and returned values
I have placed in-line comments to clarify the intent of any hard-to-understand passages of my code
I have updated the README to cover introduced code changes
I have added tests that prove my fix is effective or that my feature works
I have given the PR a name that clearly describes the change, written in imperative form (context).
I have requested a reviewer and an assignee (assignee is responsible for merging). This applies only if you have write access to the repo, otherwise feel free to tag a maintainer to add a reviewer and assignee.

Checklist for reviewers

Each PR comes with its own improvements and flaws. The reviewer should check the following:

the code is readable
the code is well tested
the code is documented (including return types and parameters)
the code is easy to maintain

Author checklist after completed review

I have added a line to the CHANGELOG describing this change, in a section
reflecting type of change (add section where missing):
- added: when you have added new functionality
- changed: when default behaviour of the code has been changed
- fixes: when your contribution fixes a bug
- maintenance: when your contribution is relates to repo maintenance, e.g. CI/CD or documentation

Checklist for assignee

PR is up to date with the base branch
the tests pass
(if the PR is not just maintenance/bugfix) the PR is assigned to the next milestone. If it is not, propose it for a future milestone.
author has added an entry to the changelog (and designated the change as added, changed, fixed or maintenance)

Co-authored-by: Hurricane <hurricane@hermes.local>

GiGiKoneti

Appreciate the fast fix and the regression tests for checkpoint loading.
cc: @joeloskarsson

joeloskarsson

Thanks, a few comments.

joeloskarsson · 2026-06-29T16:41:55Z

+            "neural_lam.train_model.torch.load",
+            return_value={"hyper_parameters": {"args": args}},
+        ),
+        patch("neural_lam.train_model.MODELS", {"hi_lam": DummyPredictor}),


Can we test this also for graph_lam, as that would have caught the original crash? (now that load_forecaster_module_from_checkpoint actually uses all args).

Related to this — neither new test constructs an args that is missing the *_gnn_type attributes entirely. That is the actual "older checkpoint" scenario from #686 (checkpoints saved before these CLI flags existed), and the getattr(..., "InteractionNet") fallback path in build_predictor/load_forecaster_module_from_checkpoint is currently untested. Could we add a case for that alongside the graph_lam case you asked for above?

joeloskarsson · 2026-06-29T16:42:18Z

        output_clamping_upper: dict[str, float] | None = None,
        g2m_gnn_type: str = "InteractionNet",
        m2g_gnn_type: str = "InteractionNet",
+        **_kwargs: object,


To be consistent with rest of codebase

Suggested change

**_kwargs: object,

**kwargs: object,

joeloskarsson

Oops, hit the wrong box.

sadamov

I would avoid **kwargs on the constructors here. It swallows typos in programmatic use, removes the very error that just caught this, and contradicts requirement #1 of #672 (explicit APIs). It also doesn't fix the loader.

Instead one small build_predictor helper that adds the hierarchical kwargs only for hierarchical models. Signatures stay explicit, both call sites de-duplicated, both bugs fixed. Good enough until #672 lands, without running the risk of having lingering **kwargs in the codebase. What do you think @joeloskarsson @gitcommit90?

Make the kwargs explicit like this:

def build_predictor(predictor_class, args, config, datastore):
    kwargs = dict(datastore=datastore, graph_name=args.graph, ...,
                  g2m_gnn_type=args.g2m_gnn_type, m2g_gnn_type=args.m2g_gnn_type)
    if issubclass(predictor_class, BaseHiGraphModel):
        kwargs["mesh_up_gnn_type"] = args.mesh_up_gnn_type
        kwargs["mesh_down_gnn_type"] = args.mesh_down_gnn_type
    return predictor_class(**kwargs)

sadamov · 2026-06-30T08:27:15Z

I am also noting that this was a clear gap in our test-suite that we should definitely fill with this PR.

joeloskarsson · 2026-07-01T10:27:49Z

To me the specialized instantiation for different model classes doesn't feel so maintainable, but I can also see the potential issues with the kwargs approach. Overall the way we instantiate models in train_model.py is no good, but this should probably be solved with #672.

For a fix right now I don't have strong opinions, so we can go with something like @sadamov's suggestion.

- Route training and checkpoint reload through build_predictor - Omit hierarchical mesh GNN kwargs for graph_lam - Remove GraphLAM **_kwargs swallow - Add regression test for graph_lam kwargs

Sir-Sloth-The-Lazy · 2026-07-06T16:54:17Z

+        g2m_gnn_type=getattr(args, "g2m_gnn_type", "InteractionNet"),
+        m2g_gnn_type=getattr(args, "m2g_gnn_type", "InteractionNet"),
+    )
+    if getattr(args, "model", None) in ("hi_lam", "hi_lam_parallel"):


This works today since MODELS only has three entries, but @sadamov's proposal above was issubclass(predictor_class, BaseHiGraphModel), not a hardcoded model-name check. As written, adding a new hierarchical model means remembering to add its name to this tuple nothing fails loudly if it's missed, it just silently falls back to InteractionNet for the hierarchical only kwargs. Worth switching to the issubclass check to match what was actually discussed?

sadamov

build_predictor reads well, just the one thing we already agreed on: gate on issubclass(predictor_class, BaseHiGraphModel) rather than the model-name tuple so future hierarchical models are covered automatically, suggestion inline.

sadamov · 2026-07-13T11:30:24Z

+        g2m_gnn_type=getattr(args, "g2m_gnn_type", "InteractionNet"),
+        m2g_gnn_type=getattr(args, "m2g_gnn_type", "InteractionNet"),
+    )
+    if getattr(args, "model", None) in ("hi_lam", "hi_lam_parallel"):


Gate on the class hierarchy so any future BaseHiGraphModel subclass is covered without editing this list.

Suggested change

if getattr(args, "model", None) in ("hi_lam", "hi_lam_parallel"):

if issubclass(predictor_class, BaseHiGraphModel):

This needs BaseHiGraphModel adding to the from .models import ... line at the top (line 22, outside this diff so I cannot suggest it inline).

Use issubclass(predictor_class, BaseHiGraphModel) instead of a model-name tuple so future hierarchical models get mesh_up/down kwargs automatically (sadamov review on mllam#688). Co-authored-by: Hermes Agent

sadamov · 2026-07-15T07:06:13Z

@gitcommit90 #673 is waiting on this fix here, are you still working on #688, otherwise I can also take over for this last small adjustment and merge.

@gitcommit90

Pull in @gitcommit90's fix for the graph_lam GNN kwarg bug (mllam#688, closes mllam#686): train_model passed hierarchical-only mesh_up/mesh_down gnn_type kwargs to every model, so `--model graph_lam` crashed with a TypeError. mllam#688 adds a `build_predictor` helper that only forwards those kwargs to hi_lam / hi_lam_parallel. This is carried here only to give the notebook a working training run. mllam#688 is the canonical fix - merge it first, then rebase/merge mllam#577 on top. Conflict resolved in tests/test_train_model_warnings.py: kept both main's test_steps_to_log_validation and mllam#688's new gnn-kwarg tests.

Fix graph LAM GNN kwarg handling

4a24927

Co-authored-by: Hurricane <hurricane@hermes.local>

GiGiKoneti approved these changes Jun 28, 2026

View reviewed changes

sadamov requested a review from joeloskarsson June 29, 2026 08:03

joeloskarsson approved these changes Jun 29, 2026

View reviewed changes

joeloskarsson requested changes Jun 29, 2026

View reviewed changes

sadamov self-requested a review June 30, 2026 08:23

sadamov reviewed Jun 30, 2026

View reviewed changes

joeloskarsson mentioned this pull request Jul 1, 2026

RFC: a general probabilistic forecasting interface #685

Open

refactor: add build_predictor helper for explicit GNN kwargs (mllam#686)

640b795

- Route training and checkpoint reload through build_predictor - Omit hierarchical mesh GNN kwargs for graph_lam - Remove GraphLAM **_kwargs swallow - Add regression test for graph_lam kwargs

Sir-Sloth-The-Lazy reviewed Jul 6, 2026

View reviewed changes

sadamov reviewed Jul 13, 2026

View reviewed changes

Merge branch 'main' into fix/issue-686-gnn-kwargs

718deaf

This was referenced Jul 13, 2026

Enforce 100% type-hint coverage and align typing to PEP 585/604 #673

Open

Replace mypy type-check gate with ty (astral-sh) GiGiKoneti/neural-lam#1

Merged

refactor: gate hierarchical GNN kwargs on BaseHiGraphModel

e705585

Use issubclass(predictor_class, BaseHiGraphModel) instead of a model-name tuple so future hierarchical models get mesh_up/down kwargs automatically (sadamov review on mllam#688). Co-authored-by: Hermes Agent

sadamov mentioned this pull request Jul 14, 2026

Dead code: load_forecaster_module_from_checkpoint is defined but never called #703

Open

sadamov requested a review from joeloskarsson July 15, 2026 07:08

	if getattr(args, "model", None) in ("hi_lam", "hi_lam_parallel"):
	if issubclass(predictor_class, BaseHiGraphModel):

Uh oh!

Conversation

gitcommit90 commented Jun 27, 2026

Describe your changes

Issue Link

Type of change

Checklist before requesting a review

Checklist for reviewers

Author checklist after completed review

Checklist for assignee

Uh oh!

GiGiKoneti left a comment • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

joeloskarsson left a comment

Choose a reason for hiding this comment

Uh oh!

joeloskarsson Jun 29, 2026

Choose a reason for hiding this comment

Uh oh!

Sir-Sloth-The-Lazy Jul 6, 2026

Choose a reason for hiding this comment

Uh oh!

joeloskarsson Jun 29, 2026

Choose a reason for hiding this comment

Uh oh!

Uh oh!

joeloskarsson left a comment

Choose a reason for hiding this comment

Uh oh!

sadamov left a comment

Choose a reason for hiding this comment

Uh oh!

sadamov commented Jun 30, 2026

Uh oh!

joeloskarsson commented Jul 1, 2026

Uh oh!

Sir-Sloth-The-Lazy Jul 6, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

sadamov left a comment

Choose a reason for hiding this comment

Uh oh!

sadamov Jul 13, 2026

Choose a reason for hiding this comment

Uh oh!

sadamov commented Jul 15, 2026

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

5 participants

GiGiKoneti left a comment •

edited

Loading

Sir-Sloth-The-Lazy Jul 6, 2026 •

edited

Loading