tasks/task2b.md

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

The current ROOT-GNN environment setup is incomplete.

Running:

```bash
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:

```text
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:

```toml
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:

```toml
[[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:

```toml
[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:

```bash
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:

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

10. Validate imports explicitly:

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

11. Run:

```bash
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:

```python
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:

```text
chore: fix ROOT-GNN PyTorch DGL environment
```