Skip to content

Fix load_custom_model when base class package is unavailable#5441

Merged
rtimms merged 4 commits intomainfrom
fix/load-custom-model-missing-base-class
Apr 10, 2026
Merged

Fix load_custom_model when base class package is unavailable#5441
rtimms merged 4 commits intomainfrom
fix/load-custom-model-missing-base-class

Conversation

@rtimms
Copy link
Copy Markdown
Contributor

@rtimms rtimms commented Apr 9, 2026

Summary

  • Serialise.load_custom_model now falls back to pybamm.BaseModel (with a warning) when the recorded base_class can't be re-imported, instead of raising ImportError. This restores the pre-Model to json #5389 contract and makes models authored in third-party packages loadable by anyone, since the model content is stored fully symbolically.
  • Consolidated the fallback logic: the "builtins.object" sentinel now lives in the top-level guard alongside "" and "pybamm.BaseModel", removing a dead branch inside the except clause (importing builtins never raises).
  • Updated test_import_base_class_non_builtin_object to assert the new warn-and-fallback behaviour, and simplified its setup to round-trip through save_custom_model + a single JSON field edit rather than hand-rolling the model JSON.

Context

Reported by a user loading a HalfCellGITTModel defined in package on a machine without that package installed.

Before #5389, load_custom_model(filename, battery_model=None) defaulted to pybamm.BaseModel() and never tried to import anything. #5389 started storing f"{cls.__module__}.{cls.__name__}" and importlib.import_module-ing it at load time, with no fallback for non-pybamm subclasses. This PR keeps the opportunistic re-import (so subclass behaviour still round-trips when the package is installed) but degrades gracefully when it isn't.

Test plan

  • pytest tests/unit/test_serialisation/test_serialisation.py -k "custom_model or base_class" — 8 passed
  • CI

🤖 Generated with Claude Code

Restore the pre-#5389 contract where `Serialise.load_custom_model` always
produced a plain `BaseModel` when the recorded subclass isn't available in
the loading environment. Previously it raised `ImportError`, making models
authored in third-party packages unloadable by anyone without that package
installed, even though the model is stored fully symbolically.

Now the loader opportunistically re-imports the original class when
possible (preserving subclass-specific Python behaviour) and otherwise
warns and falls back to `BaseModel`.

Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
@rtimms rtimms requested a review from a team as a code owner April 9, 2026 20:51
Covers the common case of a custom model that subclasses (or is reparented
to) the lithium-ion BaseModel: after save + load, the recorded base class
is re-imported, the options are restored as BatteryModelOptions, and the
default_geometry / default_spatial_methods still cover the particle
domains — so Simulation can discretise without the caller passing any
extra geometry/submesh/spatial_methods dicts.

Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
@codecov
Copy link
Copy Markdown

codecov Bot commented Apr 9, 2026

Codecov Report

✅ All modified and coverable lines are covered by tests.
✅ Project coverage is 98.31%. Comparing base (d46beb3) to head (c70873e).
⚠️ Report is 1 commits behind head on main.

Additional details and impacted files
@@           Coverage Diff           @@
##             main    #5441   +/-   ##
=======================================
  Coverage   98.31%   98.31%           
=======================================
  Files         336      336           
  Lines       30700    30699    -1     
=======================================
  Hits        30183    30183           
+ Misses        517      516    -1     

☔ View full report in Codecov by Sentry.
📢 Have feedback on the report? Share it here.

🚀 New features to boost your workflow:
  • ❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.

Copy link
Copy Markdown
Member

@BradyPlanden BradyPlanden left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Looks good, could you also add a changelog entry

rtimms and others added 2 commits April 10, 2026 09:25
Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
@rtimms rtimms merged commit 428a181 into main Apr 10, 2026
25 of 26 checks passed
@rtimms rtimms deleted the fix/load-custom-model-missing-base-class branch April 10, 2026 13:38
rtimms added a commit that referenced this pull request Apr 10, 2026
* Fall back to BaseModel when custom model base class can't be imported

Restore the pre-#5389 contract where `Serialise.load_custom_model` always
produced a plain `BaseModel` when the recorded subclass isn't available in
the loading environment. Previously it raised `ImportError`, making models
authored in third-party packages unloadable by anyone without that package
installed, even though the model is stored fully symbolically.

Now the loader opportunistically re-imports the original class when
possible (preserving subclass-specific Python behaviour) and otherwise
warns and falls back to `BaseModel`.

Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>

* Test that li-ion BaseModel round-trips with options and defaults intact

Covers the common case of a custom model that subclasses (or is reparented
to) the lithium-ion BaseModel: after save + load, the recorded base class
is re-imported, the options are restored as BatteryModelOptions, and the
default_geometry / default_spatial_methods still cover the particle
domains — so Simulation can discretise without the caller passing any
extra geometry/submesh/spatial_methods dicts.

Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>

* Remove verbose inline comments

Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>

* Update changelog for #5441

Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>

---------

Co-authored-by: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
(cherry picked from commit 428a181)
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants