Hugging Face's logo

HWresearch
/
GNN4Colliders

PyTorch
gnn4colliders
graph-neural-network
high-energy-physics
particle-physics
event-classification
Model card Files
xet
 Community
18
GNN4Colliders / tasks /task7.md
ho22joshua's picture
ho22joshua
rewriting codebase (#7)
916755e
preview code
|
Raw History Blame
19.9 kB
# Task 7: Implement Metadata-Aware Dataset Caching, Splits, and Batching
Implement the dataset-orchestration layer for the active ROOT-GNN workflow.
This task builds on:
```text
Task 4: shared physics feature construction
Task 5: ROOT-GNN graph construction
Task 6: ROOT/Awkward ingestion and architecture-neutral samples
```
Before making changes, read:
```text
AGENTS.md
docs/architecture.md
docs/migration.md
```
Then inspect:
```text
src/gnn4colliders/data/
src/gnn4colliders/features/
src/gnn4colliders/graphs/
tests/unit/
tests/parity/
tests/integration/
legacy/root_gnn_dgl/
```
Focus on active legacy behavior around:
```text
RootDataset / LazyDataset
PreBatchedDataset
fold_selection
GraphDataLoader
padding
cache loading
```
Do not implement the model or training lifecycle yet.
---
# Goal
Implement the layer that turns processed event samples and graphs into deterministic train/validation/test datasets and batches ready for a future training loop.
The intended flow is:
```text
ROOT files
->
EventSample
->
shared physics features
->
graph construction
->
GraphSample
->
cache
->
metadata-based split selection
->
batching / optional padding
->
GraphBatch / DataLoader
```
The future training layer should consume this API without needing to understand legacy dataset internals.
The design should also remain compatible with future model families such as `root_transformer`.
---
# 1. Replace legacy tracking columns with named metadata
The new implementation must not preserve the legacy positional tracking tensor as its public representation.
Legacy behavior currently encodes:
```text
tracking[:, 0] -> fold identifier
tracking[:, 1] -> event weight
```
Replace that with explicit named metadata.
At minimum:
```python
metadata["fold"]
metadata["weight"]
```
The representation may be typed rather than a raw dictionary, but downstream users must not need to know positional column numbers.
Parity should preserve semantics, not the legacy representation.
---
# 2. Use a typed EventMetadata structure
Prefer a small typed structure rather than an unconstrained dictionary.
For example:
```python
from dataclasses import dataclass, field
from typing import Any, Mapping
@dataclass(frozen=True)
class EventMetadata:
fold: int
weight: float
sample_id: str
extra: Mapping[str, Any] = field(default_factory=dict)
```
The exact implementation may differ if existing Task 6 code suggests a better design.
Requirements:
* `fold` has explicit meaning
* `weight` has explicit meaning
* every sample has a stable `sample_id`
* optional future metadata goes into `extra`
* avoid hidden positional semantics
* keep the structure lightweight and serializable
Do not build a large metadata framework.
---
# 3. Stable sample identity
Every event/sample should have a stable identity.
Use a deterministic identifier based on source provenance, for example conceptually:
```text
source file + tree + entry index
```
The exact representation may be a string or structured fields.
Example:
```python
sample_id = "file.root:Events:12345"
```
The same source event should receive the same `sample_id` across repeated runs.
Use this ID for:
* debugging
* cache round-trips
* batch alignment tests
* future inference output association
Do not use a random UUID.
---
# 4. Separate event metadata from cache metadata
Use distinct concepts and names.
## EventMetadata
Describes one physics event.
Examples:
```text
fold
weight
sample_id
event_number
run_number
source_file
```
## CacheMetadata
Describes a cache artifact or cache chunk.
Examples:
```text
feature schema version
graph schema version
cache schema version
input source
preprocessing configuration
software version
```
Do not use one generic `metadata` object for both purposes.
---
# 5. Explicit schema versioning
Introduce simple schema-version constants.
Prefer something like:
```python
FEATURE_SCHEMA_VERSION = 1
GRAPH_SCHEMA_VERSION = 1
CACHE_SCHEMA_VERSION = 1
```
Use clear names and keep versioning simple.
Cache metadata should include these versions.
The purpose is to prevent stale caches from silently loading after preprocessing or graph semantics change.
Do not implement semantic-version machinery or a migration framework.
---
# 6. EventSample contract
If Task 6 currently exposes a sample type, update it if necessary so the new architecture is explicit.
Conceptually:
```python
@dataclass(frozen=True)
class EventSample:
objects: ...
label: ...
global_features: ...
metadata: EventMetadata
```
Avoid exposing legacy `tracking`.
If Task 6 already introduced named metadata, preserve that design and extend it.
---
# 7. GraphSample contract
Introduce or update a graph-level sample representation.
Prefer something conceptually like:
```python
@dataclass
class GraphSample:
graph: dgl.DGLGraph
label: torch.Tensor
global_features: torch.Tensor | None
metadata: EventMetadata
```
This should represent one processed event for ROOT-GNN.
Keep it small and explicit.
Do not create inheritance hierarchies.
---
# 8. GraphBatch contract
Introduce a clear batch representation.
Conceptually:
```python
@dataclass
class GraphBatch:
graph: dgl.DGLGraph
labels: torch.Tensor
global_features: torch.Tensor | None
metadata: ...
```
Batched metadata must preserve one-to-one correspondence with events.
For example:
```python
batch.metadata.fold
batch.metadata.weight
batch.metadata.sample_id
```
or equivalent.
The implementation may use tensors/lists internally, but the meaning must remain explicit.
---
# 9. Metadata batching rules
Define deterministic batching rules.
For homogeneous numeric scalar fields:
```text
fold
weight
event_number
run_number
```
prefer tensors where practical.
For strings or heterogeneous metadata:
```text
sample_id
source_file
sample_name
```
prefer ordered sequences/lists.
For example:
```python
BatchMetadata(
fold=torch.tensor([0, 1, 1]),
weight=torch.tensor([1.0, 0.8, 1.2]),
sample_id=[
"a.root:Events:1",
"a.root:Events:2",
"b.root:Events:4",
],
)
```
Do not collapse heterogeneous metadata into an opaque tensor.
---
# 10. Named-field convention
Add or enforce this project-wide rule:
> Domain meaning should be represented by named fields or schemas, not only by positional column indices.
Do not introduce new APIs that require users to know things like:
```text
column 0 means fold
column 1 means weight
column 5 means event number
```
This rule should be reflected in `AGENTS.md` if not already present.
---
# 11. Legacy compatibility mapping
At the compatibility boundary, explicitly map:
```text
legacy tracking[:, 0] -> EventMetadata.fold
legacy tracking[:, 1] -> EventMetadata.weight
```
Add parity tests for this mapping.
Do not reconstruct a legacy tracking tensor in normal production APIs.
A test-only compatibility helper is acceptable if needed.
---
# 12. Fold selection
Implement fold selection through named metadata:
```python
sample.metadata.fold
```
not positional tracking.
Provide a clean API.
Conceptually:
```python
def select_folds(
samples,
folds: Collection[int],
):
...
```
Verify:
* correct membership
* deterministic ordering
* no mutation
* sample IDs remain stable
---
# 13. Split definitions
Provide an explicit abstraction for train/validation/test splits.
A small dataclass is acceptable:
```python
@dataclass(frozen=True)
class SplitDefinition:
train_folds: frozenset[int]
validation_folds: frozenset[int]
test_folds: frozenset[int]
```
The exact API may differ.
Do not hardcode specific fold numbers in generic data code.
Keep the split definition independent of Hydra for now.
---
# 14. Keep split machinery extensible
Fold-based splitting is the current implementation, but avoid making the entire dataset layer inherently dependent on folds forever.
Structure the code so alternative selectors could be added later without rewriting batching or caching.
Do not over-engineer a generic query language.
A simple separation between:
```text
split definition
selection implementation
dataset/batching
```
is sufficient.
---
# 15. Event weights
Preserve legacy weight semantics as:
```python
sample.metadata.weight
```
Maintain:
* numeric value
* sign
* dtype where relevant
* event alignment
Do not apply absolute value here.
Loss-weight interpretation belongs to the future losses/training task.
---
# 16. Cache architecture
Implement processed-data caching with clear boundaries.
Prefer a cache abstraction around:
```text
GraphSample
```
for the active ROOT-GNN path.
However, design the cache code so it does not prevent a future representation-neutral cache layer.
Do not tightly couple all cache APIs to ROOT-GNN terminology if a generic sample cache abstraction is straightforward.
---
# 17. Consider two cache levels
Evaluate whether the architecture should support two conceptual cache stages:
```text
Level 1:
ROOT
->
architecture-neutral EventSample / normalized feature data
Level 2:
features
->
GraphSample
```
Do not necessarily implement both fully in this task.
At minimum, structure the cache layer so a future ROOT-Transformer can reuse expensive ROOT preprocessing without requiring DGL graph caches.
If implementing only graph caching now is simplest, document the future separation point.
Do not create speculative unused code just to represent Level 1.
---
# 18. CacheMetadata
Introduce a typed or explicit cache metadata representation.
Conceptually:
```python
@dataclass(frozen=True)
class CacheMetadata:
cache_schema_version: int
feature_schema_version: int
graph_schema_version: int
source_files: tuple[str, ...]
tree_name: str
```
Additional stable fields may be added if necessary.
Keep cache metadata human-readable and inspectable where practical.
---
# 19. Cache provenance
Record enough provenance to understand how cached data was produced.
Consider including:
```text
source files
tree name
feature schema version
graph schema version
cache schema version
preprocessing config fingerprint/hash
GNN4Colliders version or git commit if practical
```
Do not make Git availability a hard runtime requirement.
If recording the Git commit is awkward or unreliable, make it optional.
---
# 20. Configuration fingerprint
If preprocessing configuration affects cache validity, create a deterministic representation or fingerprint.
Avoid using unstable Python object hashes.
Prefer deterministic serialization of the relevant configuration followed by a stable hash.
Do not include unrelated training configuration.
Cache validity should reflect preprocessing/representation choices only.
---
# 21. Cache validity
On load, verify relevant schema/provenance metadata.
Do not silently load obviously incompatible caches.
Raise a clear error for incompatible schema versions.
Do not implement automatic migration of old cache formats in this task.
---
# 22. Cache format
Choose a simple format appropriate for the active DGL/PyTorch path.
DGL-native graph serialization is acceptable.
Isolate file-format details behind APIs such as:
```python
save_graph_cache(...)
load_graph_cache(...)
```
The future trainer should not depend on file layout.
---
# 23. Cache round-trip
Test:
```text
GraphSample
->
save
->
load
->
GraphSample
```
Verify:
* graph topology
* node features
* edge features
* label
* global features
* EventMetadata.fold
* EventMetadata.weight
* EventMetadata.sample_id
* optional metadata
* schema metadata
Sample IDs must survive exactly.
---
# 24. Chunking
Implement deterministic chunking only if needed for realistic dataset size.
Do not preserve legacy `np.array_split` boundaries merely for historical fidelity.
Exact chunk IDs are implementation details unless external consumers require them.
Chunk ordering must be deterministic.
---
# 25. Lazy loading
If chunked caches are implemented, allow data to be loaded without materializing the entire dataset.
Keep it simple.
Do not reproduce the legacy ring buffer unless profiling later demonstrates a need.
---
# 26. Deterministic ordering
Preserve deterministic ordering through:
```text
input
processing
cache
split
batch
```
unless explicit shuffling is requested.
Use `sample_id` in tests to verify ordering and alignment.
Do not rely on filesystem iteration order.
---
# 27. Deterministic shuffling
Use explicit local RNG state.
For example:
```python
generator = torch.Generator()
generator.manual_seed(seed)
```
Avoid:
* process-wide NumPy seed mutation
* hidden calls to `torch.manual_seed`
* unseeded shuffles
Given identical:
```text
dataset
configuration
seed
```
the sample order must be reproducible.
Verify ordering using `sample_id`.
---
# 28. Batching
Batch graphs with DGL's normal batching primitives where appropriate.
Preserve alignment between:
```text
graph
label
global features
EventMetadata
```
Tests should verify batch correspondence using `sample_id`.
---
# 29. Collation
Implement an explicit collate function.
Conceptually:
```python
def collate_graph_samples(
samples: Sequence[GraphSample],
) -> GraphBatch:
...
```
It should:
1. batch graphs
2. stack labels
3. stack globals if present
4. collate named metadata
5. preserve sample ID order
Do not hide this logic in an anonymous lambda.
---
# 30. DataLoader
Provide a small loader-construction API.
Conceptually:
```python
def build_graph_dataloader(
dataset,
*,
batch_size: int,
shuffle: bool,
seed: int,
num_workers: int = 0,
):
...
```
Keep training-specific decisions out of this function.
Do not implement DDP yet.
---
# 31. Padding
Inspect active legacy padding behavior.
Legacy modes may include:
```text
NONE
STEPS
FIXED
NODE
```
Implement:
```text
NONE
```
plus only the modes required by active standard configs.
Do not carry forward unused padding modes merely for completeness.
---
# 32. Padding semantics
For implemented padding modes, test:
* padded node count
* padded edge count
* feature shape
* dummy feature values
* sample/metadata alignment
Do not preserve arbitrary legacy constants unless they are active compatibility requirements.
If a historical fixed padding size is not required, make it configurable or defer it.
---
# 33. Keep padding separate from graph construction
Task 5 should continue to build real event topology only.
Padding belongs in batching/orchestration.
Do not modify core graph-building semantics to support padding.
---
# 34. Integration flow
Add an end-to-end integration test covering:
```text
tiny ROOT fixture
->
EventSample
->
feature builder
->
graph builder
->
GraphSample
->
cache
->
split
->
DataLoader
->
GraphBatch
```
Verify:
* `sample_id`
* fold
* weight
* label
* globals
* graph topology
* deterministic ordering
* cache round-trip
* batch alignment
Do not involve a neural network.
---
# 35. Tests
Add or update tests for:
```text
EventMetadata
sample_id generation
legacy tracking -> EventMetadata mapping
SplitDefinition
fold filtering
CacheMetadata
cache schema mismatch
cache round-trip
deterministic shuffle
metadata batching
sample alignment
padding
DataLoader output
```
Prefer small deterministic fixtures.
---
# 36. No model or training implementation
Do not implement:
```text
models/root_gnn/
losses
metrics
optimizers
training loop
checkpointing
DDP
```
Those belong to later tasks.
---
# 37. No Hydra wiring yet
Do not wire these components directly to Hydra.
Use typed Python APIs.
Hydra will configure these objects from outside in a later task.
---
# 38. No premature optimization
Do not add:
```text
numba
custom CUDA
torch.compile
large process-global caches
complex multiprocessing
```
Correctness, clarity, reproducibility, and stable interfaces come first.
Document performance concerns for later profiling.
---
# 39. Documentation
Update `docs/architecture.md` to reflect the intended contracts:
```text
EventSample
architecture-neutral event representation
EventMetadata
named per-event provenance/training metadata
GraphSample
graph-specific representation for ROOT-GNN
GraphBatch
training-ready batch
CacheMetadata
cache schema and provenance
```
Explicitly document:
```text
Legacy:
tracking[:, 0] = fold
tracking[:, 1] = weight
New:
EventMetadata.fold
EventMetadata.weight
```
Document that named fields replace positional semantics intentionally.
Update `docs/migration.md` once Task 7 reaches parity.
If `AGENTS.md` does not already include the rule, add:
```text
Prefer named schemas and typed structures over positional domain semantics.
```
Keep documentation concise.
---
# Validation
Run focused tests:
```bash
uv run pytest tests/unit/data -v
uv run pytest tests/unit/graphs -v
```
Run integration tests:
```bash
uv run pytest tests/integration -v
```
Run parity tests:
```bash
uv run pytest tests/parity -v
```
Run the full suite:
```bash
uv run pytest
```
Run lint/format checks:
```bash
uv run ruff check src tests
uv run ruff format --check src tests
```
Inspect:
```bash
git status
git diff
```
Verify:
* legacy code remains unchanged
* no positional tracking API is introduced
* sample IDs are stable
* event metadata and cache metadata are separate
* schema versions are persisted
* cache round-trip preserves metadata
* no model/training implementation was added
* no large cache artifacts are tracked
* no `/global/cfs` or `/pscratch` paths are hardcoded
* no unrelated files are included
---
# Completion criteria
Task 7 is complete when:
1. `EventMetadata` or equivalent typed named metadata exists.
2. Legacy fold semantics map to `EventMetadata.fold`.
3. Legacy weight semantics map to `EventMetadata.weight`.
4. Every event has a deterministic stable `sample_id`.
5. New APIs no longer expose positional tracking as the preferred representation.
6. `EventSample` uses named metadata.
7. `GraphSample` uses named metadata.
8. `GraphBatch` preserves named metadata and sample order.
9. `CacheMetadata` is distinct from `EventMetadata`.
10. Feature, graph, and cache schema versions are explicit.
11. Cache compatibility is validated on load.
12. Cache round-trip preserves sample IDs and event metadata.
13. Fold selection is implemented using named metadata.
14. Split definitions are explicit and not hardcoded.
15. Deterministic shuffling works with an explicit seed.
16. Graph batching preserves graph/label/metadata/global alignment.
17. Required active padding modes are implemented.
18. A DataLoader produces a clean `GraphBatch`.
19. End-to-end preprocessing-to-batch integration passes.
20. Relevant legacy parity tests pass.
21. Full tests pass.
22. No model or training implementation has been added.
23. Legacy code remains untouched.
---
# Completion report
Report:
1. EventMetadata design
2. sample ID strategy
3. legacy tracking-to-metadata mapping
4. EventSample contract
5. GraphSample contract
6. GraphBatch contract
7. metadata batching rules
8. CacheMetadata design
9. schema-version strategy
10. cache format
11. cache provenance/fingerprint strategy
12. cache round-trip results
13. split/fold API
14. deterministic shuffle behavior
15. DataLoader API
16. padding modes implemented
17. unit test results
18. integration test results
19. parity results
20. intentional deviations from legacy representation
21. performance concerns deferred
22. unresolved questions
23. validation commands and results
After validation succeeds, create one Git commit containing only Task 7 changes.
Use:
```text
feat: implement metadata-aware dataset caching and batching
```
Before committing, inspect the final diff and ensure no unrelated files are included.