Add translate named data streams and paired-by-time multi-resolution loading - #1394
Add translate named data streams and paired-by-time multi-resolution loading#1394mcgibbon wants to merge 14 commits into
Conversation
Lets several samplers over equal-length datasets be shuffled independently of one another rather than in lockstep, which the translate multi-stream loader needs for its independently-sampled pairing groups.
Defines what the objectives will declare they need from each named data stream (StreamRequirements, ObjectiveDataRequirements) and the merge the data layer consumes (TranslateDataRequirements.from_objectives): per-stream variable names as an ordered union, window length as the pointwise maximum of the objectives' schedules, and the pairing groups as connected components of the stream co-occurrence graph. Deriving pairing from co-occurrence rather than configuring it means a sampling knob cannot contradict the objectives: streams compared in some objective are sampled at the same valid times, streams that never co-occur are sampled independently.
Generalizes fme.coupled's CoupledBatchData from two fixed components to an arbitrary set of named streams: a dict[str, BatchData] with per-stream access, device movement, the epoch ace's LossSchedule consumes, and a merge that assembles independently sampled groups' batches into one. Time-window manipulation is deliberately not fanned out over all streams — the objectives apply it per stream, knowing which of theirs is input, target and forcing.
TranslateDataLoaderConfig is a list of streams, each binding an ace dataset config to the domain it serves; a stream's name defaults to its domain. There is no sampling knob: get_gridded_data builds one dataset per stream and one PairedStreamDataset per derived pairing group, so streams an objective compares are sampled at the same valid times and streams no objective compares are sampled independently. A group validates its alignment when it is built — equal timesteps, and sample start times matching index by index, with the first offending index and its two times named — because pairing by index is otherwise silently wrong. A group's length is the minimum over its streams, since objectives asking for different window lengths leave different numbers of valid start times. TranslateGriddedData zips the groups' loaders, giving each group's sampler its own seed offset: without that, two equal-length groups get the same permutation and are index-paired in fact while claiming to be independent. n_batches is the minimum over groups and n_samples counts samples per group. It also publishes dataset_info keyed by domain, the pairing key ComponentPoolConfig.build consumes, validating that streams sharing a domain describe compatible datasets.
Builds a 1°/2°/4° component pool from the loader's dataset_info and pushes a loaded 1° batch through a transform built that way, which is the one path where a stream/domain misbinding would otherwise be silent. Updates the example configs to stop saying only component_pool: is implemented.
The component_pool: blocks do not parse until the channels-on-transforms rework lands, so say so rather than calling the block implemented.
The duplicate-pairing-group check was unreachable: a stream in two groups also fails the partition check, which ran first and reported it as a partition mismatch. Reorder so each error names what is actually wrong. DataRequirements carries allow_missing_variables, but StreamConfig.get_dataset does not pass it to DatasetABC.build, so a requirements object with it set would load every variable and claim success. Refuse it instead.
Batches are moved to the device but never scattered spatially, so under spatial parallelism every co-rank would receive the whole global field instead of its shard. require_no_spatial_parallelism exists for exactly this kind of known-incorrect path; fail rather than run it.
A group's length and its start-time alignment are computed once, when the PairedStreamDataset is built, so both are only correct because a stream's number of valid start times is set by its schedule's longest window (XarrayDataset uses n_timesteps.max_value) rather than by the current epoch. Nothing was holding that; a change to epoch-dependent lengths would have made the cached length and alignment prefix silently stale.
num_data_workers buys one worker pool per group, so process count and prefetched host memory scale with the number of groups. And TranslateGriddedData's docstring claimed a stream partition it does not check; name where the guarantee actually comes from.
mcgibbon
left a comment
mcgibbon
left a comment
There was a problem hiding this comment.
Claude: pre-review pass (independent agent, no authoring context). Verdict: ready for review after the four fixes I pushed — 58c0a7a. I found no correctness bug in the pairing, alignment, merge or sizing logic; the things I changed are a silently-wrong path under spatial parallelism, two requirements that were silently ignored, one unreachable error branch, and a test locking in an invariant the design depends on. Details in the thread comments; two design questions I deliberately left for Jeremy rather than settling here.
What I verified rather than took on trust:
seed_offsetplacement and back-compat. Every existing caller (fme/ace/data_loading/getters.py:71,fme/downscaling/data/config.py, the two test modules) passes no offset and getsseed=self._seedexactly as before. The stated reason for not offsetting the epoch instead holds, and is in fact stronger than the description says:GenericDataLoader.set_epochcallsself._dataset.set_epoch(epoch), which reachesXarrayDataset._global_epoch; that feeds bothDatasetItem's epoch (→BatchData.epoch→LossSchedule.init_for_epochatfme/ace/stepper/single_module.py:1632) and_ensure_epoch_synchronized's_n_timesteps_schedule.get_value(epoch). So an epoch offset would corrupt the loss schedule and shift the window-length schedule per group. Offsetting the seed is right.- Independent-group sampling. The independence test is load-bearing: patching
seed_offset=group_indextoseed_offset=0ingetters.pymakestest_independently_sampled_streams_are_not_time_lockedfail (assert not True), everything else still passing. On the RandomSampler question — there is no such path here.Distributed.get_sampleralways returns aDistributedSampler, and translate always goes through it, so the mechanism is uniform serially and distributed (ace only reachesRandomSamplerviasample_with_replacement, which this config deliberately does not expose). That also meansset_epoch/alternate_shuffle'sisinstance(..., DistributedSampler)branches are always taken. - Pointwise-max schedule. Correct at every epoch, including 0 and epochs between milestones: each input schedule is piecewise-constant with breakpoints only at its own milestones, so the max is piecewise-constant with breakpoints only in the union — evaluating at the union epochs and holding between them is exact.
ValidatedMilestonesrejects a milestone at epoch ≤ 0, sostart_value = max(get_value(0))is not shadowed by an epoch-0 milestone. - Group length is not off by one.
XarrayDataset._n_initial_conditions = total_timesteps - n_timesteps.max_value + 1, somin(len(dataset))over a group is exactly the shared valid-start count (8 times / window 4 → 5, matching the test), and the alignment prefix compare is over equal-length slices. Alignment and timestep are validated inPairedStreamDataset.__init__, which runs before the loader exists, so no sample can be drawn misaligned. GriddedDataABCconformance. All six abstract members implemented;n_samples/n_batches/batch_sizematchfme.ace'sGriddedDatasemantics (local batch size,n_batches * batch_size).TranslateDataLoaderadds no__init__args, soGenericDataLoader.subset'sself.__class__(dataset=, sampler=, collate_fn=, **kwargs)re-instantiation is safe — andsubsetfreezes each group's sampler order independently, which keeps within-group pairing while leaving groups independent.- Tests pass (79 in
fme/translate, 479 infme/core/distributed+fme/core/test_rand.py+fme/ace/data_loading),pre-commit run --filesclean on every changed file.
Fixes pushed (each its own commit):
9ff7dbb— the duplicate-pairing-group check inTranslateDataRequirements.__post_init__was unreachable (a stream in two groups also fails the partition check, which ran first and misreported it), andDataRequirements.allow_missing_variableswas silently dropped by the loader. Reordered, and refuse the flag.84ddd70—require_no_spatial_parallelismguard inget_gridded_data.4a0c199— test that a scheduled window keeps the group's start-time count.58c0a7a—num_data_workersis per pairing group; drop a docstring guarantee that isn't checked.
Revert any of these freely if you disagree — 2 in particular is a behavior change I judged worth making, not an obvious one.
Two things I think are Jeremy's call, not mine, both flagged in threads below: whether n_batches = min over groups is the right definition of an epoch when groups are very unequal, and whether an unconsumed configured stream should error.
| num_replicas=self._distributed.total_data_parallel_ranks, | ||
| rank=self._distributed.data_parallel_rank, | ||
| seed=self._seed, | ||
| seed=self._seed + seed_offset, |
There was a problem hiding this comment.
Claude: this is the right place — the alternative (mutating the returned sampler's seed) would either duplicate the num_replicas/rank logic or reach into DistributedSampler internals, and adding a full seed= override would break set_seed's single point of control. Back-compatible for all four existing call sites.
One property worth a line of comment, because it is easy to over-claim from the docstring: DistributedSampler.__iter__ seeds with self.seed + self.epoch, so offset k at epoch e yields the same permutation as offset 0 at epoch e+k. Within any single epoch the offsets differ and the groups are genuinely independent, which is the thing that matters here (and alternate_shuffle's alternate_seed(epoch) keeps them distinct too). But the offsets are a shift along one seed axis, not independent streams, so nothing later should rely on cross-epoch independence between two groups. If you want that property outright, self._seed + seed_offset * <large stride> gets it for free.
| raise ValueError("At least one pairing group's loader is required.") | ||
| self._loaders = list(loaders) | ||
|
|
||
| def __len__(self) -> int: |
There was a problem hiding this comment.
Claude: for Jeremy, not settled in pre-review. min over groups makes the epoch as long as the shortest pairing group. The reasoning given (a shorter group re-drawn mid-epoch correlates its samples within the epoch) is right, but the cost lands squarely on the transfer-learning program, where the two streams are independent groups with no reason to be the same length: if ERA5 has 10x the samples of SHiELD, an "epoch" touches a tenth of ERA5, and everything keyed to epochs (LR schedule, n_forward_steps milestones, max_epochs, validation cadence) is now denominated in the short stream. Coverage over many epochs is fine, since each epoch reshuffles — this is about what "epoch" means, not about wasted data.
The alternatives are all worse in their own way (cycling the short group correlates it; padding it duplicates samples; per-group epoch counters make the schedule ambiguous), so I don't think min is wrong. But it deserves to be a decision you made rather than a consequence you find out about when the first transfer-learning run's LR schedule finishes 10x early. Worth at least a line in the config docstring telling the user to size the streams comparably.
| unused = sorted(configured - required) | ||
| if unused: | ||
| raise ValueError( | ||
| f"These data streams are configured but no objective consumes them: " |
There was a problem hiding this comment.
Claude: on the debatable one — I think the error is right, and for a stronger reason than "don't silently load unused data". Since pairing is derived from the objectives, a stream no objective names has no group, so the failure mode of tolerating it is not a wasted read, it's a stream that is configured, never sampled, and never mentioned again. The overwhelmingly likely cause is a typo in a stream name or an objective edited without its stream, and both are exactly what you want to fail on. It is also the easy direction to relax later: turning an error into a warning breaks nobody.
The one case to keep an eye on is PR 3 introducing a train-only objective, since train_data: and validation_data: are validated against the same requirements — a stream only a train-side objective consumes would still have to appear in validation_data:. That's a fine constraint today (both example configs mirror their stream lists) but it is worth remembering when the objective list grows a train/validation distinction.
| ) | ||
|
|
||
|
|
||
| def _domain_dataset_info( |
There was a problem hiding this comment.
Claude: two things about comparing every stream against entries[0] rather than pairwise or cumulatively.
assert_compatible_with skips a field entirely when either side is None (fme/core/dataset_info.py:100-135). So if the first stream has infer_timestep: false (timestep=None) or a NullVerticalCoordinate, that field goes unchecked against every other stream serving the domain, and streams 2 and 3 are never compared to each other — three streams can pass while two of them disagree. Then the returned DatasetInfo takes horizontal_coordinates, vertical_coordinate, spatial_mask_provider and timestep from the reference alone, so the domain can be published with timestep=None even though the streams that were compatible-by-omission all carry one.
Not reachable from either example config (both name real datasets with inferred timesteps), so I left it. If you want it airtight, compare each new stream against the accumulated info and rebuild info as you go, rather than holding the reference fixed.
| raise ValueError("StreamRequirements requires at least one variable name.") | ||
|
|
||
| @property | ||
| def n_timesteps_schedule(self) -> IntSchedule: |
There was a problem hiding this comment.
Claude: nit — this constructs a whole DataRequirements (and copies names) to read a two-branch property. IntSchedule.from_constant is right there, and the union is normalized in one place either way:
if isinstance(self.n_timesteps, IntSchedule):
return self.n_timesteps
return IntSchedule.from_constant(self.n_timesteps)The delegation reads as if it were load-bearing (the comment says it is) when what it buys is avoiding a duplicated three-line branch, at the cost of a reader chasing into fme.ace.requirements to find out that's all it does. Your call; I didn't want to churn the seam.
| ) | ||
|
|
||
| @property | ||
| def stream_names(self) -> list[str]: |
There was a problem hiding this comment.
Claude: stream_names has no caller and no test. Its natural use would have been letting TranslateGriddedData check that its group loaders' streams partition stream_properties — the guarantee its docstring claimed — but GenericDataLoader doesn't expose its dataset, so the check can't be made from there. I softened that docstring in 58c0a7a instead. Either drop this property or wait for PR 3 to actually need it; as it stands it reads like an API someone forgot to wire up.
Same observation, lower stakes, for TranslateGriddedData.variable_metadata and TranslateBatchData.to_cpu — no consumer yet. Those two have direct counterparts in fme.ace/fme.coupled that PRs 3-4 will plausibly want, so forward-declaring them is more defensible.
| streams={name: batch.to_cpu() for name, batch in self.streams.items()} | ||
| ) | ||
|
|
||
| def pin_memory(self) -> "TranslateBatchData": |
There was a problem hiding this comment.
Claude: pin_memory mutates self.streams and returns self, while to_device/to_cpu build a new TranslateBatchData. BatchData's own methods are all copy-returning, so the odd one out here is a reader trap even though torch's pin-memory hook doesn't care which you do. Making it return TranslateBatchData(streams={...}) costs nothing and makes the three read the same.
Separately, on the surface as a whole: this defines __getitem__/__contains__/__iter__/__len__ with Mapping semantics (__iter__ yields keys) without being a Mapping, so batch.keys(), batch.items() and **batch all fail while sorted(batch) and name in batch work. Either inherit collections.abc.Mapping[str, BatchData] and get the rest for free, or drop the dunders and let callers use .streams — the half-Mapping is the version a reader has to test to understand. Not blocking; the tests only use the dunders that exist.
Checking each stream against the first one only was not transitive: assert_compatible_with skips a comparison either side declines to make, so a stream with a NullVerticalCoordinate is compatible with any two streams that disagree with each other. And the timestep comparison is guarded on the caller's timestep being set, so a stream with no inferred timestep passed against one that had it — after which the domain published that None for every stream serving it, since the returned info takes its coordinates from the first entry. Comparing every ordered pair fixes both, and makes taking those shared fields from the first entry sound.
Building a DataRequirements to read a two-branch property was indirection for its own sake; IntSchedule.from_constant is the same normalization.
Neither is referenced. Wiring stream_names into a partition check in TranslateGriddedData would not enforce more than is already enforced: GenericDataLoader keeps its dataset private and wraps it in a _Subset on subset(), which forwards no such property, so the check would have to read the streams back from the caller — and the caller derives them from TranslateDataRequirements, whose __post_init__ already validates that the pairing groups partition its streams. __iter__ is likewise unused, and no other DatasetABC defines one; torch reaches the samples through __getitem__.
It had __getitem__/__iter__/__len__ but no keys()/items()/values(), so iterating yielded names while ** unpacking failed — an awkward half-mapping for the code that will use it most, PR 3's objectives. Inheriting collections.abc.Mapping supplies the rest from __getitem__ and keeps them in agreement with it. Also documents at pin_memory why it mutates in place while to_device/to_cpu copy: BatchData's and fme.coupled's do the same, and torch's DataLoader calls it on a batch it has just collated and is about to hand over.
1 participant
PR 2 of the
fme/translateseries, targeting the long-lived base branchfeature/translate-skeleton(PR #1360) rather thanmain.Translate training reads several named data streams at once — the same state at
1°/2°/4°, or ERA5 alongside SHiELD. Some of those streams must be sampled at the
same valid times (a translation objective comparing 1° and 2° fields of one state
is meaningless if the two are drawn from different times) and some must not be
(the transfer-learning program's two streams are unpaired distributions). This
adds
fme/translate/data/: named streams, that pairing, and the multi-streamGriddedDataABCimplementation training will consume.Whether two streams are sampled together is derived, never configured. Each
objective declares what it needs from each stream it consumes; every objective
ties its streams together, and the connected components of the resulting graph
are the pairing groups. Streams in a group share a time index; groups are
sampled independently. Grouping is transitive, so objectives over (a, b) and
(b, c) put all three on one index. A
sampling:knob in the config couldcontradict the objectives; a derivation cannot.
The requirements type is the seam PR 3 (objectives) will produce: it constructs
one
ObjectiveDataRequirementsper objective and callsTranslateDataRequirements.from_objectives, which merges per-stream variablenames as an ordered union and per-stream window length as the pointwise maximum
of the objectives'
IntSchedules (ace loads a window of the scheduled length andeach objective samples a prefix of it, so loading the longest window any objective
asks for serves them all). Nothing in the seam depends on the component pool: the
loader takes variable names from the requirements, not from the pool's domain
channel lists.
Alignment within a group is validated when the group is built, since pairing by
index is otherwise silently wrong: the streams must share a timestep, and their
sample start times must match index by index, with the first offending index and
its two times named in the error. A group's length is the minimum over its
streams, because objectives asking for different window lengths leave different
numbers of valid start times.
n_samples/n_batches/seeding semantics under unequal groups:n_batchesis theminimum over the groups (iteration stops with the shortest, so no group is
re-drawn mid-epoch), and
n_samplesisn_batches * batch_size— samples drawnper group, not summed across groups, so it stays the count a per-batch metric
divides by. Each group's sampler is built with its own
seed_offset; withoutthat, two equal-length groups receive the same permutation and are index-paired
in fact while claiming to be independent (a test fails if the offset is removed).
TranslateGriddedData.dataset_infois keyed by domain, not stream: that isthe pairing key
ComponentPoolConfig.buildconsumes, and streams sharing adomain are checked to describe compatible datasets via
DatasetInfo.assert_compatible_with, with their variable metadata and labelsunioned.
Changes:
fme.translate.data.requirements:StreamRequirements,ObjectiveDataRequirements, andTranslateDataRequirements.from_objectives—the objective-to-data seam and its merge (name union, pointwise-max schedule,
pairing groups).
fme.translate.data.config:StreamConfig(domain,dataset, optionalnamedefaulting to the domain viastream_name) andTranslateDataLoaderConfig, mirroringfme.ace.data_loading.config.DataLoaderConfig'sfield names and validating in
__post_init__.fme.translate.data.batch_data:TranslateBatchData, aMapping[str, BatchData]generalizingfme.coupled's two-component batchtype, plus
TranslateCollateFn.fme.translate.data.dataset:PairedStreamDataset, an N-streamGenericDatasetthat pairs by index and validates timestep and start-timealignment at build time.
fme.translate.data.dataloader,.gridded_data,.getters:TranslateDataLoader(one per pairing group, all behavior fromGenericDataLoader),TranslateGriddedData(zips the groups, publishesdomain-keyed
dataset_info), andget_gridded_data(config, requirements, train).fme.core.distributed.Distributed.get_sampler: newseed_offsetargument(default 0, existing callers unaffected) so several samplers over equal-length
datasets can be shuffled independently rather than in lockstep. Added here
rather than mutating the returned sampler so the spatial-parallel
num_replicasand
ranklogic stays in one place.fme/translate/examples/: the README table and the two configs' headers nowsay which blocks parse —
train_data:/validation_data:with this PR, andcomponent_pool:only once PR 1's remaining rework (channels-on-transformswith derived domain load lists, int-declared latent blocks, identity
normalization) lands, which is tracked separately and is why the loader takes
variable names from the requirements seam rather than the pool config.
Design points worth a reviewer's attention:
PairedStreamDatasetis translate-local rather than a reuse offme.downscaling'sFineCoarsePairedDataset, which is hard-wired to afine/coarse pair and to single-timestep items (its
BatchItemDatasetAdaptersqueezes the time dimension away) and so cannot carry the time windows a
forward-prediction objective needs.
fme.coupled.CoupledDatasetis the closeranalogue but is likewise fixed at two named components. What is reused is the
machinery under both:
GenericDataset/GenericDataLoadersupply subsetting,epoch propagation, shuffle control,
batch_sizeandn_samples.Unequal but integer-ratio timesteps within a group (
fme.coupled'sn_steps_fastindex scaling) is a deliberate non-feature; the docstring namesit as the path if an objective ever needs it.
Batches are moved to the device but not spatially scattered, following
fme.coupled: spatial model parallelism would need each stream scatteredagainst its own grid, which no translate module supports yet. Data-parallel
sharding is handled.
get_gridded_datacallsDistributed.require_no_spatial_parallelism, so a spatially-parallel runfails rather than handing every co-rank the whole global field.
DataRequirements.allow_missing_variablesis refused rather than ignored:StreamConfig.get_datasetdoes not plumb it toDatasetABC.build, so arequirements object carrying it would load every variable and report success.
ace's
time_bufferwindow-reuse options are not offered: they draw severaloutput batches from one preloaded window, which would break the index pairing.
A configured stream that no objective consumes is an error, not a silent
unused load.
Tests added
If dependencies changed, "deps only" image rebuilt and "latest_deps_only_image.txt" file updated