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