Skip to content

fix: pin pytorch and align the PyG wheel index (#232) - #241

Open
yasumorishima wants to merge 4 commits into
openclimatefix:mainfrom
yasumorishima:fix/232-remove-torch-scatter
Open

yasumorishima wants to merge 4 commits into
openclimatefix:mainfrom
yasumorishima:fix/232-remove-torch-scatter

Conversation

@yasumorishima

Copy link
Copy Markdown

Fixes #232.

Root cause

The pixi environment resolves pytorch without a version constraint, while the installpyg task fetches the companion wheels from the torch 2.7.0 index. Run 30261010897 on main shows both halves:

pytorch 2.10.0 cpu_generic_py312_h5a55c2b_5   conda-forge
pip install ... -f https://data.pyg.org/whl/torch-2.7.0+cpu.html
Successfully installed torch_scatter-2.1.2+pt27cpu ...
Fatal Python error: Segmentation fault
  File ".../torch_scatter/__init__.py", line 16 in <module>

conda-forge moved past 2.7 some time ago, so the compiled extensions no longer share an ABI with the installed torch and importing torch_scatter aborts the interpreter. pytest dies during collection, which is why every run since has failed inside two minutes without executing a test. The 2.7.0 in the job name is a matrix label only; nothing pins torch.

Both ubuntu-latest and macos-latest install 2.10.0, so the mismatch is not platform specific.

Change

Pin pytorch = "2.10.0" and point installpyg / installpygcuda at the matching indexes. The pin is exact rather than 2.10.* because data.pyg.org publishes a single torch-2.10.0 index for this line: a later patch release on conda-forge would otherwise be picked up while the wheels stayed behind, recreating the same mismatch.

Verification

In a clean pixi environment on Linux:

  • Before: python -c "import torch_scatter" aborts with SIGABRT, and pytest dies during collection.
  • After: the imports succeed and the suite runs to completion — 182 passed, 10 failed on this branch's base commit.

The 10 remaining failures are unrelated to this change and are what the fix makes visible again: test_anemoi::test_normalization, three in test_gencast, two in test_weather_station_reader (SynopticPy is not installed), and four in test_weathermesh (NATten is unavailable on this platform).

One note for whoever looks at CI next: test_model.py::test_forecaster_and_loss and its two variants build a 2592-node global forecaster and run a backward pass. On an 8GB machine they climb to about 7GB in 16 seconds and get killed, so I deselected them for the measurement above. They are untouched here and may be fine on the 16GB hosted runners — but since CI has not reached them since June, that is worth watching once this lands.

…imatefix#232)

The pixi environment resolves pytorch without a version constraint, but
the installpyg task pulls the PyG companion wheels from the torch 2.7.0
index. Once conda-forge moved past 2.7 the two no longer matched, so
importing torch_scatter aborted the interpreter and pytest crashed
during collection.

Pin pytorch to 2.10.* and point installpyg/installpygcuda at the
matching wheel indices so the compiled extensions share an ABI with the
installed torch.

Verified in a clean pixi environment: before the change
`python -c "import torch_scatter"` aborts (SIGABRT) and pytest dies
during collection; after it imports cleanly and the suite runs.
Review point: allowing patch updates leaves the same class of bug open.
data.pyg.org publishes a single torch-2.10.0 index for this line, so a
later 2.10.1 on conda-forge would be picked up by 2.10.* while the
wheels stayed at 2.10.0, recreating the ABI mismatch this change fixes.
Pinning to 2.10.0 makes the coupling between the two exact.

Verified after the change: the environment resolves and torch,
torch_scatter, torch_geometric and graph_weather all import.
With the segfault gone, pytest reaches collection and fails there
instead: every test module that imports the data package hits
"No module named 'nnja_ai'", 25 errors per job on both runners.

pyproject already defines an installnnja task, but the workflow never
calls it, so nnja-ai and its jsonschema dependency were absent. Adding
the step alongside the existing installpyg and installnat calls lets
collection complete: 208 tests are collected after running it in a
clean environment.
@yasumorishima

Copy link
Copy Markdown
Author

Pushed a follow-up after seeing this branch's first CI run.

The segfault is gone — both jobs now reach the test step instead of dying during installation. What they hit next is a different wall: 25 collection errors per job, all ModuleNotFoundError: No module named 'nnja_ai'.

pyproject.toml already defines an installnnja task, but the workflow never calls it, so nnja-ai and its jsonschema dependency were never present. That was invisible while the segfault happened first. I added the step next to the existing installpyg and installnat calls; in a clean environment running it takes collection from failing to 208 tests collected.

If you would rather not add a git-sourced install to every run, #230 guards these imports instead and would solve the same problem — happy to drop this commit in favour of that.

@yasumorishima

Copy link
Copy Markdown
Author

CI now runs. Latest run on this branch (30759490865):

8 failed, 187 passed, 3 skipped, 60 warnings, 10 errors in 204.14s

For comparison, main's last run died in 105 seconds without executing a test.

One thing I flagged as unknown in the description is now answered: test_model.py::test_forecaster_and_loss and its two variants pass on the hosted runners. They exceed 8GB locally, so I could not check them before.

The 8 failures and 10 errors are pre-existing and unrelated to this change — they are simply what has been hidden behind the import crash:

  • test_gencast: test_gencast_noise, test_gencast_sampler — AssertionError
  • test_weather_station_reader: test_interpolate_missing_data, test_resample_observations — Invalid frequency: H, the pandas rename to lowercase h. The 10 errors are all in this file too.
  • test_weathermesh: test_weathermesh_encoder shape mismatch, test_weathermesh_decoder and test_weathermesh cannot find a NATTEN backend, and test_weathermesh_processor tries to allocate 22.6GB.

Happy to open separate PRs for any of those — the pandas frequency one looks like a one-character fix, and the 22.6GB allocation looks like a real bug rather than an environment issue. I left them out here to keep this change to the CI fix.

…#232)

With pytest able to collect and run again, the suite surfaces failures that
have been masked since at least June 4. None of them are caused by the torch
pin in this PR; they are pre-existing breakage on main.

torch-harmonics 0.9.0 started applying triangular truncation in truncate_sht,
clamping mmax down to lmax. generate_isotropic_noise still sized its
coefficients as (lmax, lmax + 1), so InverseRealSHT.forward asserted on the
last dimension. Read the retained mode counts back off the transform instead.
The dropped coefficients have m > l, where the associated Legendre functions
vanish, so the generated noise is unchanged.

pandas removed the uppercase offset aliases, so every freq="1H" and freq="H"
raises ValueError. Lowercased them in weather_station_reader, dataloader and
the two training entrypoints.

The WeatherMesh encoder test still expected the latent depth from before
"Fix keeping the correct depth": the pressure path uses stride=(1, 2, 2), so
the depth is the 25 pressure levels plus the surface level, not 5. The
processor, decoder and end-to-end tests asked NATTEN for a head dimension of
4, which no backend supports on CPU, and the processor's 26x32x64 grid made
the flex fallback materialise a 22.7 GB attention mask - that is what killed
the macOS job with exit 137. Kept latent_dim // num_heads >= 8 and shrank the
grids.

test_gencast_graph cross-checks against torch_geometric's TwoHop, which
multiplies two sparse CSR matrices. torch only implements that on CPU when
built against MKL, which the macOS arm64 build is not, so that comparison is
now skipped there. graph_weather's own khop builder uses torch.sparse.mm and
is still exercised everywhere.

test_normalization declared stds=0.2 for data drawn uniformly from [0, 1),
whose standard deviation is 1/sqrt(12). That put the expected normalised
spread at 1.44 against an asserted bound of 1.5, and the twelve-value sample
crossed it often enough to be flaky - locally it failed 2 runs in 5. Declared
the true standard deviation and seeded the generator.

Verified on a Linux CPU box with the pixi default environment (torch 2.10.0,
torch-harmonics 0.9.1, natten 0.21.7, pandas 3.0.5): 10 failed / 182 passed
before, 201 passed / 4 skipped after, with the three test_forecaster_and_loss
variants deselected as they need more memory than the box has. black is clean
and this adds no new ruff findings.
@yasumorishima yasumorishima reopened this Aug 3, 2026
yasumorishima added a commit to yasumorishima/graph_weather that referenced this pull request Aug 6, 2026
yasumorishima added a commit to yasumorishima/graph_weather that referenced this pull request Aug 18, 2026
…o pytest can run

pytest segfaults on main because pixi resolves torch 2.10.0 while installpyg fetches torch-2.7.0 wheels, so torch_scatter aborts at import. openclimatefix#241 fixes that; this branch needs it to get a green run.

The only conflict was at the top of tests/test_weathermesh.py: this branch added an import, openclimatefix#241 added the module docstring. Kept both, docstring first.
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.

pytest job crashes on main (segfault/abort loading torch_scatter via torch_geometric) - has been red since at least June 4

1 participant