rewriting codebase (#7)

916755e about 2 months ago

preview code

Raw History Blame Contribute Delete

12 kB

Task 2: Environment and Developer Tooling Setup

Set up a modern, reproducible Python development environment for GNN4Colliders.

This task is about:

dependency management
packaging/runtime dependencies
development tooling
linting/formatting
tests
reproducibility of the developer environment

Do not begin migrating legacy ML implementation in this task.

Before making changes, read:

AGENTS.md
docs/architecture.md
docs/migration.md
pyproject.toml
README.md

Also inspect the current Git status and existing project structure.

Goals

After this task, a developer should be able to clone the repository and run something close to:

uv sync

uv run pytest
uv run ruff check .
uv run ruff format --check .

The environment should be simple, reproducible, and appropriate for future PyTorch/DGL development.

Use uv as the Python package/environment manager.

Use Ruff for linting and formatting.

Use pytest for testing.

Do not introduce unnecessary tooling or runtime optimizations yet.

1. Use uv for project dependency management

Configure the repository to use uv.

The repository should contain a generated:

uv.lock

Use the existing pyproject.toml rather than introducing a separate requirements file unless there is a compelling compatibility reason.

Do not add Poetry, Pipenv, Conda metadata, or duplicate package-management systems.

The intended common workflow should be:

uv sync
uv run pytest
uv run ruff check .
uv run ruff format --check .

Document this workflow in the README.

2. Python version

Choose a modern Python version compatible with the scientific stack we expect to use.

Prefer Python 3.12 unless current dependency compatibility clearly requires otherwise.

Set the supported Python version explicitly in pyproject.toml.

For the ROOT-GNN environment, use requires-python = ">=3.12,<3.13" and add a .python-version file containing 3.12.

Do not inherit the legacy Python 3.8 requirement.

If PyTorch/DGL compatibility prevents using Python 3.12, report the issue rather than silently changing the target.

3. Runtime dependencies

Add the clearly required non-GPU-sensitive scientific dependencies:

numpy
awkward
uproot
scikit-learn
hydra-core

These are expected to support:

ROOT/Awkward data access
numerical operations
metrics
experiment configuration

Do not add speculative dependencies.

4. PyTorch and DGL

The project will require:

torch
dgl

However, GPU/CUDA compatibility must be treated carefully.

Before selecting versions:

inspect any relevant current environment/module information already present in the repository
inspect the legacy environment only as a reference
determine whether a straightforward modern PyTorch/DGL installation can be represented safely in the project dependency configuration

Do not blindly preserve legacy versions such as PyTorch 2.0.1 or DGL 1.1.1.

Do not hardcode CUDA-specific wheels unless there is a clear reason and the chosen approach is documented.

If DGL/PyTorch/CUDA installation requires platform-specific handling, structure the project so that:

normal developer dependencies remain clean
Perlmutter/HPC installation guidance can be documented separately
the package metadata does not unnecessarily encode one machine's CUDA stack

If there is material uncertainty about the correct DGL/PyTorch version combination, document the issue and defer the exact GPU installation choice rather than guessing.

Do not add PyTorch Geometric.

Compatibility handoff to Task 2B

When implementing both tasks from a clean checkout, use this verified Linux/Perlmutter tuple rather than resolving the latest packages:

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

DGL is not available as a usable Linux wheel from PyPI. The verified official source is the flat index below; pin the DGL extra as dgl==2.4.0+cu121 and PyTorch as torch==2.2.2, sourced from the official CUDA 12.1 PyTorch index. Constrain NumPy to <2 for this PyTorch build, and include matplotlib in the development dependencies because the legacy parity import path requires it.

[[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" }

5. Development dependency group

Create a development dependency group using the modern pyproject.toml dependency-group mechanism if supported by the selected tooling.

Include:

pytest
pytest-cov
ruff

Do not add Black, Flake8, isort, or similar tools whose role is already covered by Ruff.

Do not add Nox yet.

Do not add Numba yet.

Do not add Lightning, Kedro, DeepSpeed, CuPy, or other architectural/performance frameworks.

6. Ruff configuration

Configure Ruff in pyproject.toml.

Use Ruff for both:

linting
formatting

Adopt sensible defaults appropriate for a scientific Python package.

At minimum, enable checks covering:

Pyflakes-style errors
pycodestyle errors
import sorting

Do not turn on an extremely aggressive rule set that produces large amounts of low-value churn.

Prefer maintainability and signal over maximum lint strictness.

Configure an appropriate line length, preferably around 88–100 characters.

Do not rewrite legacy code solely to satisfy Ruff.

Exclude or appropriately scope linting so the frozen legacy/ tree does not become part of the new-code lint burden unless explicitly intended.

The new codebase should be linted.

7. Pytest configuration

Configure pytest in pyproject.toml.

Set sensible test discovery for:

tests/

Ensure running:

uv run pytest

does not accidentally collect tests from:

legacy/

unless parity tests intentionally invoke legacy behavior.

The project currently has:

tests/unit/
tests/integration/
tests/parity/
tests/fixtures/

Preserve that structure.

Add only the minimum smoke test necessary to verify the new package imports correctly, if no tests currently exist.

For example, an import/package test is acceptable.

Do not implement ML behavior in tests during this task.

8. Git ignore cleanup

Review .gitignore.

Ensure generated development artifacts are ignored, including at least:

__pycache__/
*.py[cod]
*.egg-info/
.pytest_cache/
.ruff_cache/
.coverage
htmlcov/
.venv/

Also ensure local environment artifacts and generated Python packaging metadata are not tracked.

Inspect whether any of these artifacts are currently tracked.

If generated artifacts such as:

src/gnn4colliders.egg-info/
__pycache__/

are tracked, remove them from Git tracking without altering useful source files.

Do not delete legacy model artifacts merely because they are large; the legacy tree is out of scope unless explicitly addressed.

9. Pre-commit / prek

Do not make Git hooks a hard dependency for this task.

However, evaluate whether a small hook configuration would be useful.

If adding one, keep it limited to fast checks such as:

ruff check
ruff format --check

Do not run the full test suite on every commit.

Prefer prek only if it integrates cleanly with the current project and does not create additional complexity.

If there is no clear benefit yet, defer hook tooling and document it as a future option.

10. Do not add Nox yet

Do not introduce Nox in this task.

The current development commands are simple enough to be expressed directly as:

uv run pytest
uv run ruff check .
uv run ruff format --check .

Nox may be introduced later if the project needs:

multi-version Python testing
CPU/GPU test environments
packaging validation
documentation builds
more complex CI orchestration

11. Do not add runtime optimization dependencies

Do not add optimization packages merely because they may improve speed later.

Specifically, do not add:

numba
cupy
deepspeed
lightning
torch-geometric

or similar dependencies in this task.

Runtime performance will be addressed after behavioral parity exists and profiling identifies bottlenecks.

Performance-critical areas may eventually include:

ROOT/Awkward ingestion
feature construction
graph construction
batching
DGL loading
PyTorch training
distributed training

Do not optimize them yet.

12. README update

Update the README with a concise developer setup section.

Include:

git clone ...
cd GNN4Colliders

uv sync

uv run pytest
uv run ruff check .
uv run ruff format --check .

If GPU-specific PyTorch/DGL setup requires additional steps, explain that clearly and separately.

On Perlmutter, uv's default home-directory cache can exhaust quota during large PyTorch/DGL installs. If extraction fails with Disk quota exceeded, retry with a scratch-backed cache:

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

Do not imply that training or inference is implemented yet.

Document the distinction between:

general developer environment
future GPU/HPC runtime environment

13. AGENTS.md update

If necessary, add a short environment/tooling section to AGENTS.md.

It should establish these conventions:

uv      -> dependency/environment management
ruff    -> linting and formatting
pytest  -> testing

Future agents should use:

uv run ...

for project commands where practical.

Agents should not add new dependencies casually.

Any new major dependency should have a clear architectural or runtime justification.

Avoid duplicating extensive instructions already present in AGENTS.md.

14. Optional convenience commands

Do not introduce Makefiles, task runners, or shell wrappers just to alias commands such as:

uv run pytest
uv run ruff check .

Keep the tooling stack minimal.

If there is a compelling reason for a task runner, report it instead of adding one automatically.

Validation

After making changes, validate the environment.

Run the relevant commands, preferably:

uv sync
uv run python -c "import gnn4colliders; print(gnn4colliders.__file__)"
uv run pytest
uv run ruff check .
uv run ruff format --check .

Also inspect:

git status
git diff

Ensure:

generated .egg-info is not tracked
__pycache__ is not tracked
legacy code has not been reformatted or otherwise modified
no ML implementation has been added
the lockfile is present if dependency resolution succeeded

If PyTorch/DGL cannot be installed because of platform/CUDA constraints, do not conceal the failure.

For ROOT-GNN validation, a skipped DGL test is not success. The current parity characterization suite also contains a separate expectation issue: its two-node no-self-loop edge test expects four edges, while the legacy graph constructor produces two. Report that failure without changing legacy code or weakening the test in this environment task.

Report:

what succeeded
what failed
the relevant compatibility issue
the recommended next step

Git

After validation succeeds, create a single Git commit containing only Task 2 changes.

Use a concise commit message such as:

chore: set up development environment and tooling

Before committing:

inspect git diff
ensure no unrelated files are included
ensure generated artifacts are excluded
ensure relevant validation passes

Do not amend or rewrite previous commits.

Completion report

At the end, report:

files created
files modified
runtime dependencies added
development dependencies added
Python version selected
PyTorch/DGL handling decision
validation commands run and results
anything intentionally deferred
commit hash and commit message

Keep this task limited to environment and developer tooling setup.