Parallelize file metadata inspection in 2-stage reader

[Baptiste-Arnould]

Summary

This PR parallelizes file metadata inspection during the initialization of the 2-stage reader.

Previously, metadata for all input files were inspected sequentially by rank 0. For datasets containing a large number of files, this can result in significant initialization time.

The metadata inspection is now distributed across the available reader ranks. The resulting metadata are gathered on rank 0 and reordered according to the original file order before constructing _file_time, _shape, _is_real, and _files_size.

The data-reading algorithm itself is unchanged.

Changes

• Add a helper function to inspect file metadata.
• Distribute metadata inspection across reader ranks.
• Gather metadata on rank 0.
• Restore the original file ordering before constructing reader metadata.

Testing

Tested with MPI on a dataset composed of 11800 files (~250TB in total).

Closes #74

[dalcinl]

Your changes are definitely an improvement and the implementation looks reasonable. There is just a minor thing I would like you to reconsider. You used a cyclic distribution to parallelize. That is OK in general, it is very convenient. But when order is important, it forces you to keep track of indices to reconstruct the original ordering, leading to slightly more convoluted code..

Look at pyspod/utils/parallel.py, note the _blockdist function. It is currently "private" (under the conventional initial underscore). However, we could rename and make it "public", and then use it in in the reader to distribute the metadata reading using a block-contiguous distribution. This way you would not need to sort to restore order, and the code may end up being simpler to read. Let me sketch the code below:

local_infos = []

# Distribute files across reader ranks
count, start = utils_par.blockdist(len(data_list), comm.size, comm.rank)
for i in range(start, start + count):
    file_name = data_list[i]
    file_shape, file_dtype, file_size = _inspect_file_metadata(file_name, variables[0])
    local_infos.append((file_name, file_shape, str(file_dtype), file_size))

# Gather metadata on rank 0
all_infos = comm.gather(local_infos, root=0)

and then you need to keep updating a few things to remove the sort call, etc.

Would you be willing to explore this approach and update the PR accordingly?

[Baptiste-Arnould]

Thanks for the suggestion. I updated the PR to use a block distribution for metadata inspection and made blockdist public. This removes the explicit file indices and sorting step and simplifies the reconstruction on rank 0.

[dalcinl]

LGTM.

@mengaldo Once tests pass (modulo the MPI failures that no one had taken care to fix 🤦 ) I think this one is good to go.

[codecov]

Codecov Report

❌ Patch coverage is 71.42857% with 14 lines in your changes missing coverage. Please review.
✅ Project coverage is 76.95%. Comparing base (19372d1) to head (1ee2f69).

Files with missing lines	Patch %	Lines
pyspod/utils/reader.py	68.88%	13 Missing and 1 partial ⚠️

Additional details and impacted files

@@            Coverage Diff             @@
##             main      #75      +/-   ##
==========================================
+ Coverage   76.84%   76.95%   +0.11%     
==========================================
  Files          16       16              
  Lines        3532     3549      +17     
  Branches      465      467       +2     
==========================================
+ Hits         2714     2731      +17     
  Misses        660      660              
  Partials      158      158              

☔ View full report in Codecov by Harness.
📢 Have feedback on the report? Share it here.

🚀 New features to boost your workflow:

• ❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.