Add translate named data streams and paired-by-time multi-resolution loading

fme/core/distributed/distributed.py

@@ -205,13 +205,27 @@ def get_sampler(
         dataset: torch.utils.data.Dataset,
         shuffle: bool,
         drop_last: bool = False,
+        seed_offset: int = 0,
     ) -> torch.utils.data.Sampler:
+        """
+        Get a sampler over ``dataset`` for this rank.
+
+        Args:
+            dataset: The dataset to sample from.
+            shuffle: Whether to shuffle the samples.
+            drop_last: Whether to drop the incomplete tail of the dataset.
+            seed_offset: Added to the distributed seed. Two samplers built with
+                different offsets produce different permutations of
+                equal-length datasets in the same epoch; leave it at 0 unless
+                sampling several datasets that must be shuffled independently
+                of one another rather than in lockstep.
+        """
         return torch.utils.data.DistributedSampler(
             dataset,
             shuffle=shuffle,
             num_replicas=self._distributed.total_data_parallel_ranks,
             rank=self._distributed.data_parallel_rank,
-            seed=self._seed,
+            seed=self._seed + seed_offset,
             drop_last=drop_last,
         )
 

fme/core/distributed/parallel_tests/test_get_sampler.py

@@ -135,3 +135,23 @@ def test_get_sampler_seed_reproducibility():
     second = list(dist.get_sampler(dataset, shuffle=True))
 
     assert first == second
+
+
+@pytest.mark.parallel
+def test_get_sampler_seed_offset_gives_independent_order():
+    """
+    Two samplers over equal-length datasets shuffle in lockstep unless they
+    are given different seed offsets.
+    """
+    dist = Distributed.get_instance()
+    n_dp = dist.total_data_parallel_ranks
+    dataset = _make_dataset(8 * n_dp)
+
+    set_seed(42)
+    baseline = list(dist.get_sampler(dataset, shuffle=True))
+    same_offset = list(dist.get_sampler(dataset, shuffle=True, seed_offset=0))
+    other_offset = list(dist.get_sampler(dataset, shuffle=True, seed_offset=1))
+
+    assert same_offset == baseline
+    assert other_offset != baseline
+    assert sorted(other_offset) == sorted(baseline)

fme/translate/__init__.py

@@ -5,7 +5,8 @@
 shared abstraction of a named pool of components (trainable transforms /
 encoders / decoders wrapped around a backbone stepper) mapping between named
 domains, used by the transfer learning and multi-resolution latent-stepping
-programs.
+programs, plus the named data streams (:mod:`fme.translate.data`) that pair with
+those domains.
 """
 
 from .components import (
@@ -14,6 +15,16 @@
     ComponentPoolConfig,
     TransformConfig,
 )
+from .data import (
+    ObjectiveDataRequirements,
+    StreamConfig,
+    StreamRequirements,
+    TranslateBatchData,
+    TranslateDataLoaderConfig,
+    TranslateDataRequirements,
+    TranslateGriddedData,
+    get_gridded_data,
+)
 from .domains import DomainConfig, LatentChannels
 from .modules import (
     InterpolateTransformConfig,
@@ -29,8 +40,16 @@
     "DomainConfig",
     "InterpolateTransformConfig",
     "LatentChannels",
+    "ObjectiveDataRequirements",
     "SameGridTransformConfig",
+    "StreamConfig",
+    "StreamRequirements",
     "TransformConfig",
     "TransformModuleConfig",
     "TransformSelector",
+    "TranslateBatchData",
+    "TranslateDataLoaderConfig",
+    "TranslateDataRequirements",
+    "TranslateGriddedData",
+    "get_gridded_data",
 ]

fme/translate/data/__init__.py

@@ -0,0 +1,43 @@
+"""Named data streams and the multi-stream loader that samples them.
+
+A *stream* is one named source of data bound to one component-pool domain (see
+:mod:`fme.translate.domains`). Several streams may serve one domain — ERA5 and
+IFS both feeding ``atmos_1deg`` — so the stream ``name`` (the handle objectives
+reference) is distinct from the ``domain`` it serves, and defaults to it.
+
+Which streams are sampled at the same valid times is *derived*, never
+configured: the objectives declare what each of them needs
+(:class:`ObjectiveDataRequirements`), and
+:meth:`TranslateDataRequirements.from_objectives` merges those into per-stream
+:class:`fme.ace.requirements.DataRequirements` plus the *pairing groups* —
+connected components of the graph in which every objective ties together the
+streams it consumes. Streams in a group are sampled at the same valid times;
+groups are sampled independently of one another. A sampling knob in the config
+could contradict the objectives; a derivation cannot.
+"""
+
+from .batch_data import TranslateBatchData, TranslateCollateFn
+from .config import StreamConfig, TranslateDataLoaderConfig
+from .dataloader import TranslateDataLoader
+from .dataset import PairedStreamDataset
+from .getters import get_gridded_data
+from .gridded_data import TranslateGriddedData
+from .requirements import (
+    ObjectiveDataRequirements,
+    StreamRequirements,
+    TranslateDataRequirements,
+)
+
+__all__ = [
+    "ObjectiveDataRequirements",
+    "PairedStreamDataset",
+    "StreamConfig",
+    "StreamRequirements",
+    "TranslateBatchData",
+    "TranslateCollateFn",
+    "TranslateDataLoader",
+    "TranslateDataLoaderConfig",
+    "TranslateDataRequirements",
+    "TranslateGriddedData",
+    "get_gridded_data",
+]

fme/translate/data/batch_data.py

@@ -0,0 +1,146 @@
+"""A batch of named data streams, and its collate function.
+
+:class:`TranslateBatchData` generalizes
+:class:`fme.coupled.data_loading.batch_data.CoupledBatchData` from two fixed
+components to an arbitrary set of named streams: where the coupled type has
+``ocean_data`` and ``atmosphere_data`` fields and fans every method out over
+both, this one holds a ``dict[str, BatchData]`` and fans out over its keys.
+
+The surface is deliberately narrower than the coupled type's — per-stream
+access, device movement, and the ``epoch`` a trainer needs. Time-window
+manipulation (``prepend``, ``get_start``, ``remove_initial_condition``) is done
+per stream by the objectives, which know which of their streams is the input,
+the target, and the forcing; hoisting those onto a whole-batch fan-out would
+apply them to streams that should not receive them.
+"""
+
+import dataclasses
+from collections.abc import Iterator, Mapping, Sequence
+
+from fme.ace.data_loading.batch_data import BatchData
+from fme.core.dataset.dataset import DatasetItem
+from fme.core.labels import LabelEncoding
+
+__all__ = ["TranslateBatchData", "TranslateCollateFn"]
+
+
+@dataclasses.dataclass
+class TranslateBatchData:
+    """A batch holding one :class:`BatchData` per named data stream.
+
+    Parameters:
+        streams: The batch's per-stream data, keyed by stream name.
+    """
+
+    streams: dict[str, BatchData]
+
+    def __post_init__(self):
+        if not self.streams:
+            raise ValueError("A TranslateBatchData must hold at least one stream.")
+        epochs = {name: batch.epoch for name, batch in self.streams.items()}
+        if len(set(epochs.values())) > 1:
+            raise ValueError(
+                "All streams in a batch must carry the same epoch (they are "
+                f"drawn in step by the trainer), got {epochs}."
+            )
+
+    @property
+    def epoch(self) -> int | None:
+        """The epoch every stream in this batch was drawn in.
+
+        Consumed by ace's ``LossSchedule.init_for_epoch``, which needs the epoch
+        of the data rather than of the trainer loop.
+        """
+        return next(iter(self.streams.values())).epoch
+
+    def __getitem__(self, name: str) -> BatchData:
+        return self.streams[name]
+
+    def __contains__(self, name: str) -> bool:
+        return name in self.streams
+
+    def __iter__(self) -> Iterator[str]:
+        return iter(self.streams)
+
+    def __len__(self) -> int:
+        return len(self.streams)
+
+    def to_device(self) -> "TranslateBatchData":
+        return TranslateBatchData(
+            streams={name: batch.to_device() for name, batch in self.streams.items()}
+        )
+
+    def to_cpu(self) -> "TranslateBatchData":
+        return TranslateBatchData(
+            streams={name: batch.to_cpu() for name, batch in self.streams.items()}
+        )
+
+    def pin_memory(self) -> "TranslateBatchData":
+        """Page-lock every stream's tensors; called by torch's DataLoader."""
+        self.streams = {
+            name: batch.pin_memory() for name, batch in self.streams.items()
+        }
+        return self
+
+    @classmethod
+    def merge(cls, batches: Sequence["TranslateBatchData"]) -> "TranslateBatchData":
+        """Combine batches over disjoint stream sets into one.
+
+        Used to assemble the independently-sampled pairing groups' batches into
+        the single batch a training step sees.
+        """
+        streams: dict[str, BatchData] = {}
+        for batch in batches:
+            overlap = sorted(set(batch.streams) & set(streams))
+            if overlap:
+                raise ValueError(
+                    f"Cannot merge batches sharing the streams {overlap}; each "
+                    "stream belongs to exactly one pairing group."
+                )
+            streams.update(batch.streams)
+        return cls(streams=streams)
+
+
+class TranslateCollateFn:
+    """Collates per-stream samples into a :class:`TranslateBatchData`.
+
+    One instance serves one pairing group: its keys are that group's streams,
+    and it is called with the group's paired samples (see
+    :class:`fme.translate.data.dataset.PairedStreamDataset`). Defined at module
+    level so it can be pickled to data-loader worker processes.
+    """
+
+    def __init__(
+        self,
+        horizontal_dims: Mapping[str, list[str]],
+        label_encodings: Mapping[str, LabelEncoding | None],
+    ):
+        """
+        Args:
+            horizontal_dims: Each stream's horizontal dimension names, used
+                when writing batches to netCDF.
+            label_encodings: Each stream's label encoding, or None for a stream
+                whose dataset provides no labels.
+        """
+        if set(horizontal_dims) != set(label_encodings):
+            raise ValueError(
+                "horizontal_dims and label_encodings must cover the same "
+                f"streams, got {sorted(horizontal_dims)} and "
+                f"{sorted(label_encodings)}."
+            )
+        self.horizontal_dims = dict(horizontal_dims)
+        self.label_encodings = dict(label_encodings)
+
+    def __call__(
+        self, samples: Sequence[Mapping[str, DatasetItem]]
+    ) -> TranslateBatchData:
+        return TranslateBatchData(
+            streams={
+                name: BatchData.from_sample_tuples(
+                    [sample[name] for sample in samples],
+                    horizontal_dims=dims,
+                    label_encoding=self.label_encodings[name],
+                )
+                for name, dims in self.horizontal_dims.items()
+            }
+        )