[ci][deps] Source torch-scatter from Astral's GPU index and drop the dead data.pyg.org find-links

[elliot-barn]

Description

data.pyg.org — the only host that ever served prebuilt wheels for the PyTorch Geometric extensions — has been a dangling CNAME since 2026-09-02 (pyg-team/pyg-lib#719). Every ML/data/rllib/doc/windows depset declared it as --find-links, and uv fails eagerly on an unreadable find-links URL, so raydepsets: compile all dependencies has been red on master (postmerge 19501, 19507) and was made soft_fail in #65880. The same URL is line 2 of every affected lock, so any cache-miss rebuild of the mlbuild/mlgpubuild/mllightning1gpubuild images (and the core/rllib/doc GPU images) fails at uv pip install -r <lock> too.

This PR narrows the dependency on data.pyg.org to the one package that needs it, sources torch-scatter from Astral's index, and re-enables the hard failure:

• torch-scatter now comes from Astral's GPU indexes (https://wheels.astral.sh/simple/{cpu,cu128}/). These are real PEP 503/691 indexes, so the requirement files use --extra-index-url instead of a flat find-links page. Versions encode the torch minor: 2.1.2+cpu.torch.2.9 / 2.1.2+cu.12.8.torch.2.9. Astral publishes manylinux wheels only, so the CPU pin carries ; sys_platform == 'linux'; nothing in Ray imports torch_scatter on any platform.
• torch-sparse stays, still from data.pyg.org (which came back on 2026-09-03; pyg-lib#719 closed). It is required: torch_geometric.loader.NeighborSampler.__init__ builds a SparseTensor from torch_geometric.typing, whose stub raises ImportError: 'SparseTensor' requires 'torch-sparse' when the package is absent — verified empirically on torch 2.9.0 + torch-geometric 2.5.3 — and python/ray/train/examples/pytorch_geometric/distributed_sage_example.py (CI test distributed_sage_example, train_v2_gpu) uses it. torch_sparse also imports torch_scatter in 8 modules, so scatter is a hard dependency of sparse. No other wheel source for torch-sparse exists (PyPI is sdist-only; not on Astral, download.pytorch.org or pypi.nvidia.com; no GitHub release assets), so the --find-links for it remains and this PR does not remove the dependency on that host — it narrows it to one package. Pinned explicitly (+pt29cpu / +pt29cu128) so a missing wheel fails loudly instead of falling back to the sdist.
• torch-cluster and torch-spline-conv are dropped. Nothing in python/, rllib/, release/ or doc/ imports them, torch-geometric==2.5.3 lists none of the four in requires_dist, and the example's full path (FakeDataset → RandomNodeSplit → NeighborSampler → SAGEConv) trains with only scatter + sparse installed (run in a glibc-2.36 container against the exact pinned wheels).
• Both remaining extensions are gated on sys_platform == 'linux': Astral ships manylinux wheels only, torch_sparse cannot import without torch_scatter, and nothing on Windows or macOS uses PyG. Windows therefore loses the four +pt27cpu pins, their find-links, and the matching relaxed_windows_tests entries.
• deepspeed: >=0.12.3 → ===0.18.9 in train-requirements.txt (PEP 440 arbitrary equality; the compiled constraint follows as deepspeed==0.18.9 → deepspeed===0.18.9). Why: Astral's cu128 index also publishes deepspeed 0.18.9+cu.12.8.torch.2.9 (prebuilt CUDA ops). Under PEP 440 a bare ==0.18.9 matches that local version, and --index-strategy unsafe-best-match prefers it, so without this every GPU depset would silently switch to Astral's deepspeed build and pull in nvidia-ml-py (verified with uv pip compile). === is the only specifier that excludes local versions, so deepspeed stays on the PyPI build it uses today; uv honours it as requirement and constraint, and pip-tools writes it through to requirements_compiled.txt. The resolved version is unchanged (0.18.9 before and after). If we ever want Astral's prebuilt-ops deepspeed, the opt-in is the reverse edit, === → == — not part of this PR.
• ci/ci.sh compile_pip_dependencies: the +cpu-stripping sed now removes the whole local segment. The compiled file is the GPU depsets' constraint (via the requirements_compiled_py3.1x.txt symlinks), and only a bare torch-scatter==2.1.2 lets +cu.12.8.torch.2.9 satisfy it; the old [^\b]*cpu pattern would have left 2.1.2.torch.2.9.
• .buildkite/dependencies.rayci.yml: soft_fail and its comment removed — the job hard-fails again.
• Also: dead --find-links removed from release/ray_release/byod/requirements_ml_byod_3.10.txt (it pinned none of the four); docs and stale comments updated.

Regenerated: python/requirements_compiled.txt (./ci/ci.sh compile_pip_dependencies, py3.11 + pip-tools 7.4.1) and all 32 affected locks under python/deplocks/ (bazel run //ci/raydepsets:raydepsets -- build --all-configs). Package-level delta in every lock is only the PyG lines; deepspeed stays 0.18.9, no new packages. python/requirements_compiled_py3.14.txt has no generator (it pins torch 2.7.0, which has no cp314 wheels) and has always been maintained by mirroring the hunks applied to requirements_compiled.txt (#65616, #65046, #64056 all changed identical pin counts in both), so the same header/pin deltas were applied to it by hand; its own torch-2.7.0+cpu find-links is kept to match its torch pin. That line is only ever read by pip: requirements_compiled_py${PYTHON_VERSION}.txt is staged into the base-deps/base-slim images as /home/ray/requirements_compiled.txt (those images install from the rayimg lock with uv, not from this file) and then used as -c by ci/docker/ray-image.Dockerfile, docker/ray/Dockerfile and runtime_env_container/Dockerfile — all pip install -c. pip treats an unreachable --find-links in a constraints file as a warning (5 retries, exit 0; verified against an NXDOMAIN host), whereas uv fails eagerly (exit 2; also verified). The only uv consumer is raydepsets, whose pre-hook strips the header lines. So no py3.10–3.14 image build hard-fails on that line even if data.pyg.org goes dark again; the uv exposure is the --find-links emitted into the ML/data/rllib/doc locks, which this PR narrows to torch-sparse.

Left untouched, flagged for follow-up: release/ray_release/byod/byod_horovod*.sh (install torch-2.0.1-era PyG wheels from the same dead host — separate, older problem), ci/raydepsets/tests/test_cli.py (fixture strings only).

Related issues

Follows up #65880 (temporary soft_fail). Upstream: pyg-team/pyg-lib#719, pytorch_geometric#10793, pytorch_geometric#10716.

Additional information

Not a duplicate: no open PR touches data.pyg.org, wheels.astral.sh, or the PyG pins (searched pyg, astral, find-links, torch-scatter, data.pyg.org; only unrelated dependabot torch bumps match).

Tests run locally:

• ./ci/ci.sh compile_pip_dependencies in a py3.11 venv → exit 0; diff vs master is exactly the two header lines, the PyG pins/markers and deepspeed===
• bazel run //ci/raydepsets:raydepsets -- build --all-configs → exit 0, 107 depsets compiled; the 32 changed locks differ from master only in the PyG lines and the index/find-links header, and the torch-sparse hashes are unchanged
• bazel run //ci/raydepsets:raydepsets -- build --all-configs --check (CI's exact command) → exit 0, all 107 depsets regenerate byte-identically (including the py3.14-constrained ones against the mirrored snapshot), working tree clean afterwards
• PyG minimal-set check in python:3.10-bookworm: the example's path trains with scatter+sparse only; fails without sparse, and sparse fails without scatter
• pre-commit run --files <changed source files> → shellcheck, semgrep pass
• uv pip compile probes: torch-scatter==2.1.2+cu.12.8.torch.2.9 resolves from wheels.astral.sh/simple/cu128/; deepspeed==0.18.9 flips to 0.18.9+cu.12.8.torch.2.9 with that index present, deepspeed===0.18.9 does not.

AI assistance (Claude Code) was used for this change; every changed line was reviewed by the submitter.

🤖 Generated with Claude Code

[gemini-code-assist]

Code Review

This pull request updates PyTorch Geometric (PyG) binary dependencies across various configuration and lock files, migrating from the defunct data.pyg.org wheel host to Astral's PEP 503 indexes (wheels.astral.sh). It drops torch-sparse, torch-cluster, and torch-spline-conv due to a lack of wheel sources, keeping only torch-scatter on Linux. Additionally, deepspeed is pinned using arbitrary equality (===) to avoid matching local versions from Astral's index. Feedback on the changes suggests improving the regular expression in ci/ci.sh to robustly handle pre-release or post-release versions and to enhance portability across different sed implementations.

[gemini-code-assist]

The current regular expression ==([.0-9]+)\+ only matches digits and dots for the public version part. If a pre-release, post-release, or dev-release version (e.g., 2.10.0rc1+cpu or 2.10.0.post0+cpu) is used in the future, this pattern will fail to match the full public version correctly.

Additionally, using [+] instead of \+ is more portable across different sed implementations (such as BSD sed on macOS vs GNU sed on Linux) when running in Extended Regular Expression (-E) mode, as some engines can have quirks with escaped metacharacters.

Updating the pattern to ([A-Za-z0-9.]+) and [+] resolves both issues robustly.

Suggested change

	sed -i -E 's/==([.0-9]+)\+[A-Za-z0-9.]*cpu[A-Za-z0-9.]*/==\1/g' "python/$TARGET"
	sed -i -E 's/==([A-Za-z0-9.]+)[+][A-Za-z0-9.]*cpu[A-Za-z0-9.]*/==\1/g' "python/$TARGET"

[elliot-barn]

Taken in 7bbfe65. Verified the widened class is identical on every current line shape and a no-op on the regenerated requirements_compiled.txt; it only changes behaviour for pre/post/dev public versions, which the old [.0-9]+ left unstripped.

[elliot-barn]

no longer windows requirements

[dstrodtman]

stamp from docs

[elliot-barn]

Converting to draft: verified empirically (torch 2.9.0 + torch-geometric 2.5.3, none of the four extensions installed) that torch_geometric.loader.NeighborSampler builds a SparseTensor in __init__ and raises ImportError: 'SparseTensor' requires 'torch-sparse'. python/ray/train/examples/pytorch_geometric/distributed_sage_example.py (CI test distributed_sage_example, train_v2_gpu) uses it, so dropping torch-sparse would break that test once the ML GPU image rebuilds from these locks. My earlier claim that nothing in Ray needs the extensions missed the torch_geometric.typing import path. Also: data.pyg.org resolves again as of today (pyg-lib#719 closed). Deciding between keeping torch-sparse via the restored host, self-hosting it, or changing the example before un-drafting.

[elliot-barn]

Un-drafting. 3190406ab8 restores torch-sparse (pinned 0.6.18+pt29cpu / +pt29cu128, Linux-only, still via --find-links on data.pyg.org — that host is a flat page with no PEP 503 layout, and no other wheel source for torch-sparse exists) and mirrors the header/pin deltas into requirements_compiled_py3.14.txt. torch-scatter stays on Astral's index; torch-cluster/torch-spline-conv stay dropped — verified in a glibc-2.36 container that the SAGE example's full path (FakeDataset → RandomNodeSplit → NeighborSampler → SAGEConv) trains with scatter + sparse only, fails without sparse, and sparse fails without scatter. Locally: compile_pip_dependencies exit 0; raydepsets build --all-configs 107/107; --check exit 0 byte-identical. torch-sparse wheel hashes are unchanged vs master. Description updated accordingly.

[gemini-code-assist]

Code Review

This pull request updates PyTorch Geometric (PyG) dependency management to resolve issues with PyG's dangled wheel host. It migrates torch-scatter to Astral's simple PEP 503 indexes, restricts torch-scatter and torch-sparse to Linux, and drops unused extensions (torch-cluster and torch-spline-conv). Additionally, deepspeed is pinned using arbitrary equality (===) to prevent matching local versions from Astral's index. Feedback on the changes suggests refining the regular expression in ci/ci.sh to fully support PEP 440 local version identifiers by matching hyphens and underscores.

[gemini-code-assist]

PEP 440 local version identifiers can contain hyphens (-) and underscores (_) as separators between segments (e.g., +cpu-ubuntu or +cpu_custom). The current regular expression [A-Za-z0-9.]* only matches alphanumeric characters and dots, which would stop matching at a hyphen or underscore, leaving a trailing suffix (e.g., ==2.1.2-ubuntu or ==2.1.2_custom) and resulting in an invalid version specifier.

To make the regex fully robust and PEP 440 compliant, we should include hyphens and underscores in the character class for the local version part.

Suggested change

	sed -i -E 's/==([A-Za-z0-9.]+)[+][A-Za-z0-9.]*cpu[A-Za-z0-9.]*/==\1/g' "python/$TARGET"
	sed -i -E 's/==([A-Za-z0-9.]+)[+][A-Za-z0-9._-]*cpu[A-Za-z0-9._-]*/==\1/g' "python/$TARGET"

[elliot-barn]

Taken in 0e44f9f. Note pip-compile normalizes local separators to dots (packaging's Version str), so -/_ can't actually appear in the compiled file; the wider class is harmless and verified a no-op on both compiled files.

[elliot-barn]

Two follow-ups from review:

1. Rewrote the deepspeed paragraph in the description — it read as if the PR drops === → ==; the change is >=0.12.3 → ===0.18.9 (compiled constraint == → ===) to keep deepspeed off Astral's local build. Code and in-file comment were already right.
2. Sanity-checked the data.pyg.org line still in requirements_compiled_py3.14.txt (and, via the symlinks, in requirements_compiled.txt for 3.10–3.13). That file reaches images only as a pip -c constraint: staged by docker/base-deps and docker/base-slim (which themselves install from the rayimg lock with uv, not from it), then consumed by ci/docker/ray-image.Dockerfile, docker/ray/Dockerfile and runtime_env_container/Dockerfile via pip install -c. Verified against an NXDOMAIN host: pip retries 5× and proceeds (exit 0); uv fails eagerly (exit 2). The only uv reader is raydepsets, whose pre-hook strips header lines. So no py3.14 (or 3.10–3.13) image build hard-fails on that line during an outage; the uv-side exposure lives in the locks' own --find-links, now limited to torch-sparse.