Skip to content

fix: Add AttrType to safe globals for torch serialisation - #1308

Open
HCookie wants to merge 2 commits into
perf/load-graph-strictfrom
fix/add-attrtype-to-safe-globals
Open

fix: Add AttrType to safe globals for torch serialisation#1308
HCookie wants to merge 2 commits into
perf/load-graph-strictfrom
fix/add-attrtype-to-safe-globals

Conversation

@HCookie

@HCookie HCookie commented Aug 5, 2026

Copy link
Copy Markdown
Member

Description

Addresses failing integration tests caused by #1303 as seen in https://github.com/ecmwf/anemoi-core/actions/runs/31021131375/job/92360048174?pr=1295

New graph to replace old test graph

https://sites.ecmwf.int/ecm1947/files/Anemoi/lam-graph-rebuilt.pt

As a contributor to the Anemoi framework, please ensure that your changes include unit tests, updates to any affected dependencies and documentation, and have been tested in a parallel setting (i.e., with multiple GPUs). As a reviewer, you are also responsible for verifying these aspects and requesting changes if they are not adequately addressed. For guidelines about those please refer to https://anemoi.readthedocs.io/en/latest/

By opening this pull request, I affirm that all authors agree to the Contributor License Agreement.

Addresses failing integration tests
@HCookie
HCookie requested a review from cathalobrien August 5, 2026 16:21
@HCookie HCookie self-assigned this Aug 5, 2026
@github-project-automation github-project-automation Bot moved this to To be triaged in Anemoi-dev Aug 5, 2026
@github-actions github-actions Bot added the graphs label Aug 5, 2026
@HCookie HCookie added the ATS Approval Not Needed No approval needed by ATS label Aug 5, 2026
@cathalobrien

Copy link
Copy Markdown
Contributor

thanks for the quick fix! once these pass I'm happy to merge.

@HCookie HCookie moved this from To be triaged to Under Review in Anemoi-dev Aug 5, 2026
@HCookie

HCookie commented Aug 5, 2026

Copy link
Copy Markdown
Member Author

This is failing with

_pickle.UnpicklingError: Weights only load failed. In PyTorch 2.6, we changed the default value of the `weights_only` argument in `torch.load` from `False` to `True`. Re-running `torch.load` with `weights_only` set to `False` will likely succeed, but it can result in arbitrary code execution. Do it only if you got the file from a trusted source.
Please file an issue with the following so that we can make `weights_only=True` compatible with your use case: WeightsUnpickler error:

Can only SETITEMS for dict, collections.OrderedDict, collections.Counter, but got <class 'collections.defaultdict'>

Check the documentation of torch.load to learn more about types accepted by default with weights_only https://pytorch.org/docs/stable/generated/torch.load.html.

Without a change to pytorch, this is not fixable.
It is possible that this is a problem particular to this checkpoint as checkpoints I have locally made with a recent version of anemoi-graphs do not raise this issue.

PyTorch Geometric populates BaseStorage._cached_attr as a defaultdict(set)
while building a graph. The strict weights_only=True unpickler used by
load_graph_from_file cannot deserialize a defaultdict (it only allows
SETITEMS on dict, OrderedDict and Counter), so a saved graph fails to load.
Normalise the cache to a plain dict in GraphCreator.save() before torch.save
to keep the save/load round-trip working without weakening serialization
safety.
mc4117 pushed a commit that referenced this pull request Aug 6, 2026
…emoi-graphs" (#1311)

Reverts #1303 due to issues identified in the attempt
of a hotfix in #1308

TLDR:
Some LAM graphs come with a `_cache_attr` property which is a
`defaultdict`, it is not needed, but is included by `torch_geometric`
causing issues with `torch.load`

> [!TIP]
> Integration tests passed
> https://github.com/ecmwf/anemoi-core/actions/runs/31090338148
@HCookie
HCookie changed the base branch from main to perf/load-graph-strict August 6, 2026 10:46
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

ATS Approval Not Needed No approval needed by ATS graphs

Projects

Status: Under Review

Development

Successfully merging this pull request may close these issues.

2 participants