feat: make Hydra configs installable

docs/configuration.md

@@ -0,0 +1,94 @@
+# Configuration guide
+
+The CLI composes the installed `gnn4colliders.configs/config.yaml` with semantic
+Hydra groups. A config
+describes experiment intent; it does not contain arbitrary Python module or
+class import paths.
+
+## Groups
+
+| Group | Purpose |
+| --- | --- |
+| `data` | ROOT source or graph-cache path, batch settings, and fold splits |
+| `model` | Model family and supported ROOT-GNN constructor settings |
+| `task` | Binary or multiclass loss, score, and metric semantics |
+| `trainer` | Device, seed, epochs, optimizer, scheduler, and early stopping |
+| `checkpoint` | Output directory, resume checkpoint, or pretrained checkpoint |
+| `inference` | Checkpoint, split, output path, and output format |
+| `environment` | Device and experiment output root |
+| `distributed` | Single-process/DDP selection and process-group backend |
+
+The active model groups are `root_gnn/edge_network` and
+`root_gnn/fine_tuned_edge_network`. The active task groups are
+`pretraining_multiclass`, `binary_classification`, and `tth_cp_finetune`.
+
+## Preparation
+
+`prepare` requires `data.files`, `data.cache.path`, `data.feature_branches`,
+`data.object_types`, and `data.scales`. `feature_branches` follows the shared
+seven-column feature contract: one branch/constant specification per output
+column and one entry per configured object type. `CALC_E` and `NODE_TYPE` are
+reserved derived specifications. `object_types` entries are `vector` or
+`single`. Preparation reads the configured tree in file order and writes a
+versioned `GraphSampleCache`.
+
+Example overrides are easiest to maintain in a YAML file for real datasets:
+
+```yaml
+# project-local example: data/my_events.yaml
+files: [data/events.root]
+tree_name: Events
+cache:
+  path: cache/events.pt
+feature_branches:
+  - [jet_pt]
+  - [jet_eta]
+  - [jet_phi]
+  - CALC_E
+  - [1.0]
+  - [0.0]
+  - NODE_TYPE
+object_types: [vector]
+scales: [1, 1, 1, 1, 1, 1, 1]
+fold_var: eventNumber
+weight_var: weight
+```
+
+Then compose it with `data=my_events`. The cache stores processed graph
+samples, labels, globals, named metadata, and feature/graph/cache schema
+versions. Changing the feature or graph schema requires a new compatible cache;
+loading a mismatched schema raises an error.
+
+## Common overrides
+
+```bash
+uv run gnn4colliders train \
+  data.cache.path=cache/events.pt \
+  data.batch_size=64 \
+  trainer.max_epochs=50 \
+  trainer.seed=123 \
+  environment.output_root=outputs/my_run
+```
+
+`data.batch_size` is per process. `data.splits.train_folds`,
+`validation_folds`, and `test_folds` define conventional train/validation/test
+selection and must be disjoint. Model/task mismatches are rejected during
+config validation; binary tasks require `model.out_size=1`, while multiclass
+tasks require `model.out_size=task.num_classes`.
+
+## Checkpoints and resolved configuration
+
+Training writes `epoch_####.pt` and the fully resolved configuration at
+`<environment.output_root>/resolved_config.yaml`. A checkpoint includes schema
+version, model weights/config, task config, trainer/optimizer/scheduler state,
+early stopping state, metadata, and optional RNG state. Set
+`checkpoint.resume=/path/to/epoch_####.pt` to continue a run. Set
+`checkpoint.pretrained=/path/to/epoch_####.pt` with the fine-tuned model group
+to load weights into a new task head; these options are mutually exclusive.
+
+## Environment profiles
+
+`environment=local` selects CPU by default. `environment=perlmutter` selects
+the CUDA device and a conventional output-root pattern. Profiles should hold
+device/output policy only; site-specific module loads and filesystem paths
+belong in a launcher or shell environment.

src/gnn4colliders/cli/__init__.py

@@ -3,15 +3,25 @@
 from __future__ import annotations
 
 import logging
+import os
 import sys
-from pathlib import Path
+from importlib import resources
 
 from hydra import compose, initialize_config_dir
 from omegaconf import OmegaConf
 
-from gnn4colliders.config.application import predict_or_evaluate, prepare, train
+from gnn4colliders.config.application import (
+    export_onnx,
+    predict_or_evaluate,
+    prepare,
+    train,
+)
 
-_COMMANDS = {"prepare", "train", "evaluate", "predict"}
+_COMMANDS = {"prepare", "train", "evaluate", "predict", "export"}
+
+
+def _is_main_process() -> bool:
+    return int(os.environ.get("RANK", "0")) == 0
 
 
 def _help(command: str | None = None) -> None:
@@ -19,14 +29,18 @@ def _help(command: str | None = None) -> None:
         print(f"Usage: gnn4colliders {command} [key=value ...]")
         print("Compose Hydra configuration with semantic overrides.")
     else:
-        print("Usage: gnn4colliders <prepare|train|evaluate|predict> [key=value ...]")
+        print(
+            "Usage: gnn4colliders <prepare|train|evaluate|predict|export> "
+            "[key=value ...]"
+        )
         print("Use gnn4colliders <command> --help for command help.")
 
 
 def _config(overrides: list[str]):
-    config_dir = Path(__file__).resolve().parents[3] / "configs"
-    with initialize_config_dir(version_base=None, config_dir=str(config_dir)):
-        return compose(config_name="config", overrides=overrides)
+    config_resource = resources.files("gnn4colliders.configs")
+    with resources.as_file(config_resource) as config_dir:
+        with initialize_config_dir(version_base=None, config_dir=str(config_dir)):
+            return compose(config_name="config", overrides=overrides)
 
 
 def main(argv: list[str] | None = None) -> int:
@@ -52,16 +66,22 @@ def main(argv: list[str] | None = None) -> int:
         if command == "prepare":
             print(f"prepared cache: {prepare(config)}")
         elif command == "train":
-            print(f"saved checkpoint: {train(config)}")
+            if _is_main_process():
+                print(f"saved checkpoint: {train(config)}")
+        elif command == "export":
+            if _is_main_process():
+                print(f"wrote ONNX model: {export_onnx(config)}")
         elif command == "evaluate":
             result = predict_or_evaluate(config, evaluate=True)
-            print(OmegaConf.to_yaml(result.metrics))
+            if _is_main_process():
+                print(OmegaConf.to_yaml(result.metrics))
         else:
             result = predict_or_evaluate(config)
-            print(
-                f"wrote predictions: {config.inference.output} "
-                f"({len(result.sample_ids)} events)"
-            )
+            if _is_main_process():
+                print(
+                    f"wrote predictions: {config.inference.output} "
+                    f"({len(result.sample_ids)} events)"
+                )
         return 0
     except (ValueError, FileNotFoundError, KeyError, ImportError) as error:
         print(f"gnn4colliders: {error}", file=sys.stderr)

src/gnn4colliders/configs/__init__.py

@@ -0,0 +1 @@
+"""Canonical Hydra configuration tree shipped with the package."""

{configs → src/gnn4colliders/configs}/config.yaml

@@ -6,6 +6,8 @@ defaults:
   - checkpoint: default
   - inference: default
   - environment: local
+  - distributed: single
+  - export: onnx
   - _self_
 experiment:
   name: pretraining_multiclass

{configs → src/gnn4colliders/configs}/data/local_test.yaml

@@ -1,5 +1,11 @@
 files: []
 tree_name: Events
+feature_branches: null
+object_types: null
+scales: null
+global_features: []
+fold_var: eventNumber
+weight_var: null
 batch_size: 32
 num_workers: 0
 shuffle: true

src/gnn4colliders/configs/distributed/ddp.yaml

@@ -0,0 +1,2 @@
+enabled: true
+backend: null

src/gnn4colliders/configs/distributed/single.yaml

@@ -0,0 +1,2 @@
+enabled: false
+backend: null

src/gnn4colliders/configs/export/onnx.yaml

@@ -0,0 +1,6 @@
+format: onnx
+checkpoint: null
+output: ${environment.output_root}/export/model.onnx
+opset: 17
+split: train
+overwrite: false

{configs → src/gnn4colliders/configs}/task/tth_cp_finetune.yaml

@@ -1,3 +1,3 @@
 defaults:
-  - /task/binary_classification
+  - binary_classification
   - _self_

{configs → src/gnn4colliders/configs}/trainer/debug.yaml

@@ -1,4 +1,4 @@
 defaults:
-  - /trainer/default
+  - default
   - _self_
 max_epochs: 1