GeoT optimization 2/4: Datapipes producer/consumer refactor + stream overlap

[coreyjadams]

PhysicsNeMo Pull Request

This PR is stacked 🥞 on #1741. It requires the torch SDF implementation first.. PR #1741 is merged, so this is ready to go.

This PR is a rebuild of datapipe components, particularly at the IO level. It's key changes are the following:

• There is a clear and logical separation now between "threads do IO fetching" and "streams do GPU preprocessing".
• There is an additional dataset type for "streaming" type datasets that are unbounded, so imaging "online simulation" datasets. There is an example included to show this behavior.
• The unified exteral aero recipe is updated accordingly.

What does it mean to separate IO threads from preprocessing Steams?

The PR introduces the concept of an "IO Pump": it's self priming, so to speak. It takes a few args (queue depth, worker fn), and runs a sentinal thread to keep the queue topped up. Every time an IO payload is consumed, the pump automatically queues another, ensuring the IO is consistently ready to go. The key here is the IO threads can't block the GIL, so async IO is key.

The IOPump is FIFO. Exceptions are captured and held until attempted consumption of data: if a preload a a batch fails, it doesn't actually crash everything until we go to consume that batch.

While this is based purely on threads today, in principle it's about 80% of the way to enabling processPool concurrency for prefetching too. That will break memory pinning, but for applications that simply cannot release the GIL, it could become a viable option.

For best performance, we should make sure the queue depth is larger than the batch size (so an entire batch is always ready). And if IO latency is greater than model step time, we should preload enough batches to compensate.

GPU Preprocessing with Streams

IO payloads are "consumed" in an N-1 fashion: while batch N is processed, batch N-1 is consumed on the GPU for preprocessing, then collated. The datasets use two functions to isolate work: _load_host and _consume. _load_host is what is orchestrated by the IOPump, and is CPU only and threaded. _consume, on the other hand, is single threaded (really, the MAIN thread) and handles the queuing of transforms onto the appropriate stream (an arg to _consume) as well as a deferred sync.

Cuda Stream Orchestration

Handling the thread and stream orchestration takes a little care. Python threading is relatively easy to queue and block on: we keep the queue primed by filling it to depth, making sure it has enough work so that data is always ready on the host. During training / inference, at step N, we eagerly pull batch N+1 from the dataset and start the transfer and transformations, and collation. We apply a gating to batch N (which, in this model, was pulled during N-1 processing) with cuda Events to ensure it completes before proceeding - this prevents the model from consuming it before it's ready. During the processing of batch N+1, the same gating on batch N+1 is applied right after the processing of N+2 is triggered. In this way, we keep preprocessing flowing concurrently with the model computation.

Any data preprocessing that includes blocking syncs (nonzero, HtD blocking, etc) will cause this model to stumble and introduce bubbles in the GPU flow.

RNG Updates

This PR makes some RNG updates to ensure we can have reproducibility regardless of ordering. The IO pump is FIFO but in general we shouldn't necessarily assume FIFO, and RNG reproducibility should be stable under multi-threading.

Iterable Style Dataset

There was a request to support "Simulation only" type datasets, since we have this nice prefetching model now. I included this here, and a tutorial / example as well. It's like a pytorch "iterable" dataset (so no `len`` attribute) but doing on-the-fly generation on the GPU, in side streams, concurrent with model processing. Limited currently to one-step-ahead processing.

Misc Updates

• Fixes blocking transfers in radius search Mesh ID (@loliverhennigh to review Please?)
• Fix some small GPU syncs in the unified recipe
• Updated tests

Doc Updates

I am still review the doc updates and readme updates. The PR should not merge without those.

Next Steps

I'd like to get this reviewed now, and think about some follow on PRs:

• I want to include an optional sentinel to watch for GPU syncs in the datapipe so they are easy to catch and fix.
• I would like to build features to compile the GPU preprocessing pipelines, both with torch inductor as well as cuda graphs. This could reduce launch latency overhead a lot.
• I would eventually include ProcessPool parallelism in the datasets but that is P1.

Description

Checklist

• I am familiar with the Contributing Guidelines.
• New or existing tests cover these changes.
• The documentation is up to date with these changes.
• The CHANGELOG.md is up to date with these changes.
• An issue is linked to this pull request.
• If I am implementing a new model or modifying any existing model, I have followed the Models Implementation Coding Standards.

Dependencies

Review Process

All PRs are reviewed by the PhysicsNeMo team before merging.

Depending on which files are changed, GitHub may automatically assign a maintainer for review.

We are also testing AI-based code review tools (e.g., Greptile), which may add automated comments with a confidence score.
This score reflects the AI’s assessment of merge readiness and is not a qualitative judgment of your work, nor is
it an indication that the PR will be accepted / rejected.

AI-generated feedback should be reviewed critically for usefulness.
You are not required to respond to every AI comment, but they are intended to help both authors and reviewers.
Please react to Greptile comments with 👍 or 👎 to provide feedback on their accuracy.

[copy-pr-bot]

Auto-sync is disabled for draft pull requests in this repository. Workflows must be run manually.

Contributors can view more details about this message here.

[copy-pr-bot]

This pull request requires additional validation before any workflows can run on NVIDIA's runners.

Pull request vetters can view their responsibilities here.

Contributors can view more details about this message here.

[greptile-apps]

Greptile Summary

This PR refactors the GeoT datapipe layer around a clear producer/consumer split: an IOPump dispatcher thread keeps a bounded FIFO of host-staged samples ready while the main thread performs all device-kernel launches on dedicated preprocessing CUDA streams with a one-batch lookahead so preprocessing genuinely overlaps compute. It also adds an IterableDatasetBase path for unbounded online-simulation generators, migrates all reader subsampling RNG to a thread-safe per-(seed, epoch, index) derivation via numpy.SeedSequence, and fixes several blocking transfers in the mesh radius-search path.

• IOPump + deferred-sync lookahead: a new IOPump class drives a lazy, backpressure-bounded prefetch loop; the DataLoader's _iter_prefetch drains batch N+1 before yielding batch N and defers the compute-stream wait until after the previous batch's model work is enqueued, producing genuine GPU/IO overlap.
• IterableDatasetBase: generator-style datasets run main-thread-only with a one-item lookahead on a preprocessing stream, covering online simulation use cases where the generator itself issues CUDA work.
• Thread-safe reproducible RNG: derive_seed / spawn_generator replace shared mutable generators in all readers, making subsampling deterministic regardless of worker-thread ordering.

Important Files Changed

Filename	Overview
physicsnemo/datapipes/io_pump.py	New self-priming FIFO dispatcher with backpressure via semaphore; correct lazy pull and BATCH_BOUNDARY forwarding, but stop() silently abandons the dispatcher thread after a 5-second timeout with no log or signal.
physicsnemo/datapipes/protocols.py	Introduces HostPayload/PrefetchHandle dataclasses and a clean producer/consumer split with proper thread-locking; prefetch() has a TOCTOU window that can trigger a redundant submit for the same index.
physicsnemo/datapipes/dataloader.py	Rebuilds iteration around IOPump with correct deferred-sync lookahead (batch N+1 preprocessing overlaps batch N compute); adds IterableDatasetBase path with stream handoff; one-batch lookahead logic and event ordering are well-tested.
physicsnemo/datapipes/dataset.py	Correctly splits _load_host (worker thread, CPU only) from _consume (main thread, H2D + transforms on prep stream); defer_sync plumbing and event lifecycle look correct.
physicsnemo/datapipes/mesh_dataset.py	Mirrors Dataset's producer/consumer split for mesh data; profiler annotations added around reader, device transfer, and transforms; _consume handles both CUDA and non-CUDA targets correctly.
physicsnemo/datapipes/multi_dataset.py	Adds submit/consume/_pop_events delegation across sub-datasets so MultiDataset duck-types correctly for the DataLoader's IOPump path; event aggregation from all sub-datasets is correct.
physicsnemo/datapipes/_rng.py	Adds derive_seed and spawn_generator using numpy SeedSequence for order-independent, thread-safe per-(seed, epoch, index) RNG; enables reproducible subsampling across concurrent worker threads.
physicsnemo/datapipes/readers/base.py	Moves set_generator/set_epoch into the Reader base class; adds _index_generator that derives a per-sample generator from (base_seed, epoch, index), replacing shared mutable generators in subclasses.
physicsnemo/nn/functional/neighbors/radius_search/_warp_impl.py	Replaces blocking device tensor creation with a pinned-memory + non_blocking copy; CUDA stream ordering is correct, but allocating pinned memory on every call adds cudaHostAlloc overhead that likely outweighs the benefit for a small array of grid IDs.
examples/cfd/external_aerodynamics/unified_external_aero_recipe/src/train.py	Adds a pinned scalar buffer for async D2H loss copy to avoid a per-step GPU sync, but lacks an explicit barrier between the non_blocking copy and the subsequent .item() read from pinned CPU memory.
test/datapipes/core/test_streaming.py	Comprehensive new test suite covering IOPump laziness, BATCH_BOUNDARY reassembly, FIFO ordering, error surfaces, submit/consume with opaque descriptors, iterable dataset path, drop_last, self-batching, thread isolation, and the deferred-sync overlap invariant.

_{Reviews (1): Last reviewed commit: "Update changelog." | Re-trigger Greptile}

[greptile-apps]

Silent timeout on dispatcher shutdown

stop() joins the dispatcher thread with a hard 5-second timeout but doesn't log or signal if the thread is still alive when join() returns. If dispatch_fn blocks (e.g. executor queue full, file descriptor exhaustion), the daemon thread continues running past stop() without any indication — the caller in _iter_prefetch's finally block then calls dataset.cancel_prefetch() assuming the pump is cleanly stopped. At process exit the daemon is killed, but during a long run this can accumulate hung threads and suppress the underlying problem. A simple if self._thread.is_alive(): check after join() with a warning would make the failure visible.

[greptile-apps]

Pinned memory allocation on every radius-search call

torch.tensor(..., pin_memory=torch.cuda.is_available()) registers a new pinned (page-locked) host allocation via cudaHostAlloc on every forward pass. For a tiny array (one int64 per grid, typically a handful), the CUDA host-memory registration overhead likely dominates and outweighs the non-blocking copy benefit compared to the original single blocking cudaMemcpy. Pinning is efficient only for buffers that are reused across calls. Consider caching the pinned tensor per grid configuration or simply accepting the synchronous allocation here, where the array is small enough that the blocking copy is negligible.

Note: If this suggestion doesn't match your team's coding style, reply to this and let me know. I'll remember it for next time!

[loliverhennigh]

Fine with me, pinning memory like this makes me a bit worried but its ok

[peterdsharpe]

Reviewed parts I'm codeowner on: dataset yamls, unified recipe changes, and PhysicsNeMo-Mesh changes. All those look good to me!

[negin513]

Cool to add the iterable and map style datasets.

From my read through:

1. The map-style _iter_prefetch path has the full N-1 overlap.... This is great!!
2. But the iterable dataset, _iter_iterable, doesn't have that lookahead. So it should show the same behavior as the unoptimized workflow...

Does it makes sense to split this into two PRs?

[coreyjadams]

/ok to test 255edce

[negin513]

claude flagged this for me for this example:

Darcy2D reuses its output buffers: it keepsself.output_k / self.output_p around and does .data.copy_() into them every step, so every batch this yields is the same two tensors. With use_streams=True that's exactly the race the dataloader docstring warns about... ?

[negin513]

Potential race condition since it is outside the lock?

[negin513]

LGTM!! One thing before merge is the tutorial 5 inline comment with that example.

Otherwise, look great. Does worth do a profile on it?

[peterdsharpe]

Pushed 31b6e897 (as discussed with @coreyjadams): extends this PR's history-independent RNG scheme to the two sites it didn't cover — transforms/base.py and transforms/mesh/base.py still reseeded with manual_seed(initial_seed() + epoch), which compounds across calls, so checkpoint resume at epoch N silently drew different augmentations than a continuous run. They now capture the base seed in set_generator and reseed via _rng.derive_seed(base_seed, epoch), matching the readers' SeedSequence scheme. Also updates RNG.md/DISTRIBUTIONS.md (both documented the drifting idiom as the contract) and adds two sequential-vs-direct set_epoch equivalence tests. Full test/datapipes/transforms/ suite passes (272 passed, 5 skipped). This supersedes #1817, which I'm closing.