No special treatment for ShardingRoundRobinDispatch

🚀 The feature

MPRS currently looks specifically for ShardingRoundRobinDispatch to determine the non-replicable part of the graph that gets executed in the main process before passing work to worker processes.

I would like to have the same synchronization behavior, but without the round robin or sharding aspect. Thus I would like to have a way to declare a part of the graph as non-replicable and to be executed before any worker processes. I.e. ShardingRoundRobinDispatch shouldn't get any special treatment, and features should be introduced that would allow users to be able the replicate its behavior with their own code if they wanted to.

Motivation, pitch

I'm preloading bigger chunks of data from which then worker processes produces samples in a sharded way. Currently, each worker process however loads each chunk separately from the rest since everyone is executng the same graph up the sharding point (later on, after the preloading). This results in memory consumption to scale linearly with the number of worker processes.

Instead, i would like each worker process to access the chunks via shared memory (this should be possible with posix's fork). For this I want the preloading of the chunks to be done once in the main process and then pass the same chunk (withouth sharding or round-robin behavior) to all worker processes. This is conceptually also much simpler for me to deal with, since the preloading is threaded as well.

Alternatives

No response

Additional context

No response

Start by tracing how MPRS identifies and handles ShardingRoundRobinDispatch. Determine where a user-declared non-replicable graph section would be represented and how it should run before worker processes. Done means preloading occurs once in the main process, the resulting chunks are accessible to workers, and round-robin or sharding behavior is not required.