tasks/task2b.md

Task 2B: Fix ROOT-GNN PyTorch/DGL Environment

The current ROOT-GNN environment setup is incomplete.

Running:

uv sync --extra root-gnn

currently attempts to install DGL from PyPI and fails on Linux because the required Linux DGL wheel is distributed through DGL's own wheel repository.

As a result, DGL-dependent parity tests are skipped.

This is not acceptable for the canonical ROOT-GNN development environment.

The verified target tuple is:

Python 3.12
PyTorch 2.2.2+cu121
DGL 2.4.0+cu121
CUDA runtime 12.1

Do not continue retrying DGL 2.5.0 artifacts: they returned HTTP 403 in the target environment even though some official listings advertise them. The verified Linux CPython 3.12 wheel is DGL 2.4.0+cu121 from the official torch-2.2/cu121 repository.

Goal

Make the ROOT-GNN development environment reproducible on Linux/Perlmutter and ensure DGL-dependent parity tests actually run.

Requirements

1. Inspect the current:

   • pyproject.toml
   • uv.lock
   • README environment instructions
   • AGENTS.md
   • Perlmutter/CUDA environment information available in the repository

2. Determine an explicit compatible tuple for:

   • Python
   • PyTorch
   • CUDA runtime
   • DGL

3. Prefer Python 3.12 unless compatibility evidence requires otherwise.

Set requires-python = ">=3.12,<3.13" and add .python-version containing 3.12.

4. Do not use loose unconstrained dependencies such as:

torch >= 2.0
dgl >= 2.0

for the ROOT-GNN stack.

PyTorch and DGL compatibility must be represented intentionally.

5. Configure uv so DGL is obtained from the appropriate official DGL wheel repository rather than PyPI.

Use uv's supported custom/flat index/source configuration rather than ad-hoc manual pip install commands where practical.

The DGL repo.html page is a flat wheel listing, so format = "flat" is required. Use this exact source configuration:

[[tool.uv.index]]
name = "dgl"
url = "https://data.dgl.ai/wheels/torch-2.2/cu121/repo.html"
format = "flat"
explicit = true

[tool.uv.sources]
dgl = { index = "dgl" }

Pin dgl==2.4.0+cu121 and torch==2.2.2, sourcing PyTorch from its official CUDA 12.1 index. Also constrain numpy<2; otherwise uv can select NumPy 2.x, which produces an ABI warning with this PyTorch build. Add matplotlib to the development dependencies because the legacy parity import path requires it.

6. Keep DGL architecture-specific:

[project.optional-dependencies]
root-gnn = [
    ...
]

Do not rename this back to a generic ml extra.

7. PyTorch may remain shared if that makes sense for future model families such as ROOT-Transformer.

8. Document GPU/CUDA assumptions clearly.

Do not silently make one local machine's module setup a universal package requirement.

9. After configuration, this command must succeed in the intended ROOT-GNN environment:

uv sync --extra root-gnn

If Perlmutter home-directory extraction fails with Disk quota exceeded, use a scratch-backed uv cache and rerun the command:

export UV_CACHE_DIR=/pscratch/sd/$USER/uv-cache-gnn4colliders
uv sync --extra root-gnn

10. Validate imports explicitly:

uv run python -c "import torch; print(torch.__version__)"
uv run python -c "import dgl; print(dgl.__version__)"

11. Run:

uv run pytest tests/parity

Run CUDA availability checks on a GPU allocation rather than only on a login node. Expected import output is Torch 2.2.2+cu121, Torch CUDA 12.1, DGL 2.4.0+cu121, and an available GPU.

DGL-dependent tests must run rather than skip due to DGL being absent.

12. If DGL cannot be made compatible with the selected modern PyTorch/CUDA combination, report the exact incompatibility and choose the newest well-supported compatible tuple rather than guessing.

Test behavior

Do not treat skipped DGL tests as successful ROOT-GNN validation.

The canonical ROOT-GNN validation procedure should fail early if:

import dgl

does not work.

Generic package tests may still allow optional DGL behavior, but ROOT-GNN-specific CI/development validation must require it.

The current parity characterization suite has one known expectation issue: its two-node no-self-loop edge-feature test expects four edges, while the legacy graph constructor produces two. Do not alter legacy code or weaken that test; report it separately after confirming that DGL imports and the other parity tests execute.

Out of scope

Do not:

• rewrite model code
• rewrite preprocessing
• optimize runtime performance
• add Lightning
• add PyG
• add Numba
• modify legacy behavior

This task is only about establishing a correct PyTorch/DGL development runtime.

Completion report

Report:

• selected Python version
• selected PyTorch version
• selected DGL version
• selected CUDA target
• uv source/index configuration
• commands run
• import results
• parity test results
• any Perlmutter-specific notes

Commit as:

chore: fix ROOT-GNN PyTorch DGL environment