rewriting codebase (#7)

916755e 3 months ago

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:

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:

AGENTS.md
docs/architecture.md
docs/migration.md

Then inspect:

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:

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:

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:

tracking[:, 0] -> fold identifier
tracking[:, 1] -> event weight

Replace that with explicit named metadata.

At minimum:

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:

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:

source file + tree + entry index

The exact representation may be a string or structured fields.

Example:

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:

fold
weight
sample_id
event_number
run_number
source_file

CacheMetadata

Describes a cache artifact or cache chunk.

Examples:

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:

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:

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

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

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

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:

fold
weight
event_number
run_number

prefer tensors where practical.

For strings or heterogeneous metadata:

sample_id
source_file
sample_name

prefer ordered sequences/lists.

For example:

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:

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:

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:

sample.metadata.fold

not positional tracking.

Provide a clean API.

Conceptually:

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:

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

split definition
selection implementation
dataset/batching

is sufficient.

15. Event weights

Preserve legacy weight semantics as:

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:

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:

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:

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

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:

save_graph_cache(...)
load_graph_cache(...)

The future trainer should not depend on file layout.

23. Cache round-trip

Test:

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:

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:

generator = torch.Generator()
generator.manual_seed(seed)

Avoid:

process-wide NumPy seed mutation
hidden calls to torch.manual_seed
unseeded shuffles

Given identical:

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:

graph
label
global features
EventMetadata

Tests should verify batch correspondence using sample_id.

29. Collation

Implement an explicit collate function.

Conceptually:

def collate_graph_samples(
    samples: Sequence[GraphSample],
) -> GraphBatch:
    ...

It should:

batch graphs
stack labels
stack globals if present
collate named metadata
preserve sample ID order

Do not hide this logic in an anonymous lambda.

30. DataLoader

Provide a small loader-construction API.

Conceptually:

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:

NONE
STEPS
FIXED
NODE

Implement:

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:

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:

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:

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:

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:

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:

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:

Prefer named schemas and typed structures over positional domain semantics.

Keep documentation concise.

Validation

Run focused tests:

uv run pytest tests/unit/data -v
uv run pytest tests/unit/graphs -v

Run integration tests:

uv run pytest tests/integration -v

Run parity tests:

uv run pytest tests/parity -v

Run the full suite:

uv run pytest

Run lint/format checks:

uv run ruff check src tests
uv run ruff format --check src tests

Inspect:

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:

EventMetadata or equivalent typed named metadata exists.
Legacy fold semantics map to EventMetadata.fold.
Legacy weight semantics map to EventMetadata.weight.
Every event has a deterministic stable sample_id.
New APIs no longer expose positional tracking as the preferred representation.
EventSample uses named metadata.
GraphSample uses named metadata.
GraphBatch preserves named metadata and sample order.
CacheMetadata is distinct from EventMetadata.
Feature, graph, and cache schema versions are explicit.
Cache compatibility is validated on load.
Cache round-trip preserves sample IDs and event metadata.
Fold selection is implemented using named metadata.
Split definitions are explicit and not hardcoded.
Deterministic shuffling works with an explicit seed.
Graph batching preserves graph/label/metadata/global alignment.
Required active padding modes are implemented.
A DataLoader produces a clean GraphBatch.
End-to-end preprocessing-to-batch integration passes.
Relevant legacy parity tests pass.
Full tests pass.
No model or training implementation has been added.
Legacy code remains untouched.

Completion report

Report:

EventMetadata design
sample ID strategy
legacy tracking-to-metadata mapping
EventSample contract
GraphSample contract
GraphBatch contract
metadata batching rules
CacheMetadata design
schema-version strategy
cache format
cache provenance/fingerprint strategy
cache round-trip results
split/fold API
deterministic shuffle behavior
DataLoader API
padding modes implemented
unit test results
integration test results
parity results
intentional deviations from legacy representation
performance concerns deferred
unresolved questions
validation commands and results

After validation succeeds, create one Git commit containing only Task 7 changes.

Use:

feat: implement metadata-aware dataset caching and batching

Before committing, inspect the final diff and ensure no unrelated files are included.