Allow parallel iteration over data containers

[cphyc]

PR Summary

This allows to iterate in parallel over data containers.

The syntax is

ds = yt.load(...)
ad = ds.all_data()

# Gathering everything in parallel
my_storage = {}
for sto, chunk in ad.piter(storage=my_storage, reduction="cat"):
    sto.result = {}
    sto.result["gas", "density"] = chunk["gas", "density"]

my_storage["gas", "density"]  # exists for all processes and contains everything


# Gathering everything on root
my_storage = {}
for sto, chunk in ad.piter(storage=my_storage, reduction="cat_on_root"):
    sto.result = {}
    sto.result["gas", "density"] = chunk["gas", "density"]

my_storage["gas", "density"]  # None for all processes but root


# Reduce everything
my_storage = {}
for sto, chunk in ad.piter(storage=my_storage, reduction="sum"):
    sto.result = {}
    sto.result["gas", "density"] = chunk["gas", "density"]

my_storage["gas", "density"]  # contains the sum of all the densities

This makes #4730 obsolete.

PR Checklist

• New features are documented, with docstrings and narrative docs
• Adds a test for any bugs fixed. Adds tests for new features.

[neutrinoceros]

about type-checking errors: #5195

[Copilot]

Pull Request Overview

This PR adds parallel iteration capabilities to data containers in yt, allowing users to iterate over data chunks in parallel with optional reduction operations. The implementation introduces a new piter method on data containers that supports various reduction strategies (concatenation, sum, min, max) either on all processors or just the root processor.

• Added a piter method to YTSelectionContainer for parallel iteration over data chunks
• Enhanced the parallel_objects function with reduction operations and improved type annotations
• Implemented new reduction methods (reduce, all_concat, concat) in the communication system

Reviewed Changes

Copilot reviewed 2 out of 2 changed files in this pull request and generated 4 comments.

File	Description
yt/utilities/parallel_tools/parallel_analysis_interface.py	Enhanced parallel processing infrastructure with reduction operations and modernized ResultsStorage class
yt/data_objects/selection_objects/data_selection_objects.py	Added piter method to data containers for parallel iteration with comprehensive documentation

[Copilot]

The documentation mentions 'concat' as a reduction option, but the actual parameter value is 'cat'. This inconsistency could confuse users.

Suggested change

	- concat: the storage object will contain a flattened list of
	each results.
	- cat_on_root: same as concat, but only the root processor will
	- cat: the storage object will contain a flattened list of
	each result.
	- cat_on_root: same as cat, but only the root processor will

[Copilot]

The reduce method is called with root=0, but when the operation is not distributed, the result should be available on all processors. This inconsistency with the documented behavior could cause issues.

Suggested change

	new_storage = my_communicator.reduce(to_share, op=reduction, root=0)
	if getattr(my_communicator, "distributed", False):
	new_storage = my_communicator.reduce(to_share, op=reduction, root=0)
	else:
	new_storage = my_communicator.reduce(to_share, op=reduction)

[Copilot]

Same documentation inconsistency as in the other file - 'concat' is documented but 'cat' is the actual parameter value.

Suggested change

	- concat: the storage object will contain a flattened list of
	each results.
	- cat_on_root: same as concat, but only the root processor will
	- cat: the storage object will contain a flattened list of
	each results.
	- cat_on_root: same as cat, but only the root processor will

[brittonsmith]

This looks really great. Very concise, happy to see this merged. We can use this to refactor the loop over all halos in the HaloCatalog.

[cphyc]

I need to fix the typing issues (which I'm not convinced are real) and write the doc. Please hold before merging :)

[neutrinoceros]

Yes, mypy infers reachability in a slightly broken way that may give off false negatives, maybe this lint should actually be turned off entirely.

[cphyc]

In the meantime, I've added # ignore statements to make mypy happy.

[matthewturk]

lgtm!

[chummels]

Great addition to the codebase! Will this apply to all data objects, not just ds.all_data()? Have you checked if this works with particle-based frontends? It probably does, but I don't know. If it doesn't, it might be worth having a note in the docs or an explicit error message for particle-based frontends that this it not implemented. Lastly, it might be worth including a simple unit test to ensure that the code can run (not bother with answer testing), so that we're aware if something breaks it in the future. Anyway, I'm not trying to make a ton more work, just make sure this framework remains viable in the future.

[cphyc]

I realised I never replied to your message @chummels. I haven't tried with particle data, but there's no reason for it not to work if it supports chunking.

It also works with any selection, not only all_data().

[cphyc]

aaaaand I just found a bug - let's not merge this just yet. The issue is with RAMSES reader, where we use progressive indexing (a spatial region isn't fully indexed until it's actually read). The problem is then that different MPI tasks end up reading different regions, and so which chunks may overlap a given region depends on what has been read previously.

If you use this PR's parallel reader one after the other, you may end up in a situation where, e.g. MPI task 1 thinks it needs to read CPUs #1, #2, and #3, but MPI task 2 already read #3 and indexed it, so it know it won't overlap with the region, leaving only #1 and #2 to be read. Since the MPI tasks aren't communicating this information with each other, this will lead to undefined behavior.

[cphyc]

I fixed the bug thanks to @AnatoleStorck (see 39777f6)! The test failure is unrelated, however.