fix: pin pytorch and align the PyG wheel index (#232) - #241
yasumorishima wants to merge 4 commits into
Conversation
…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.
|
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
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. |
|
CI now runs. Latest run on this branch (30759490865): 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: 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:
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.
…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.
Fixes #232.
Root cause
The pixi environment resolves
pytorchwithout a version constraint, while theinstallpygtask fetches the companion wheels from the torch 2.7.0 index. Run 30261010897 on main shows both halves: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_scatteraborts the interpreter. pytest dies during collection, which is why every run since has failed inside two minutes without executing a test. The2.7.0in the job name is a matrix label only; nothing pins torch.Both
ubuntu-latestandmacos-latestinstall 2.10.0, so the mismatch is not platform specific.Change
Pin
pytorch = "2.10.0"and pointinstallpyg/installpygcudaat the matching indexes. The pin is exact rather than2.10.*because data.pyg.org publishes a singletorch-2.10.0index 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:
python -c "import torch_scatter"aborts with SIGABRT, and pytest dies during collection.The 10 remaining failures are unrelated to this change and are what the fix makes visible again:
test_anemoi::test_normalization, three intest_gencast, two intest_weather_station_reader(SynopticPy is not installed), and four intest_weathermesh(NATten is unavailable on this platform).One note for whoever looks at CI next:
test_model.py::test_forecaster_and_lossand 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.