tasks/task2.md

# 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:

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

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

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

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

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

```text
torch
dgl
```

However, GPU/CUDA compatibility must be treated carefully.

Before selecting versions:

1. inspect any relevant current environment/module information already present in the repository
2. inspect the legacy environment only as a reference
3. 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:

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

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

---

# 5. Development dependency group

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

Include:

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

```text
tests/
```

Ensure running:

```bash
uv run pytest
```

does not accidentally collect tests from:

```text
legacy/
```

unless parity tests intentionally invoke legacy behavior.

The project currently has:

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

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

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

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

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

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

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

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

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

Future agents should use:

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

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

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

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

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

1. files created
2. files modified
3. runtime dependencies added
4. development dependencies added
5. Python version selected
6. PyTorch/DGL handling decision
7. validation commands run and results
8. anything intentionally deferred
9. commit hash and commit message

Keep this task limited to environment and developer tooling setup.