Skip to content

Parallelize file metadata inspection in 2-stage reader #75

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
Baptiste-Arnould wants to merge 2 commits into
base: main
Choose a base branch
from
Open

Parallelize file metadata inspection in 2-stage reader #75

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
3ff6ccb
Parallelize file metadata inspection in 2-stage reader
Aug 27, 2026
1ee2f69
Use block distribution for metadata inspection
Aug 30, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
10 changes: 5 additions & 5 deletions pyspod/utils/parallel.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -73,7 +73,7 @@ def create_subcomm(comm):
# import sys
# sys.stdout.flush()
# N = 7
# n, s = _blockdist(N, dims[0], coords[0])
# n, s = blockdist(N, dims[0], coords[0])
# for i in range(s, s+n):
# val = subcomm.allreduce(rank)
# subcomm.Free()
Expand All @@ -93,7 +93,7 @@ def distribute(data, comm):
shape = data.shape
index = [np.s_[:]] * len(shape)
N = shape[max_axis]
n, s = _blockdist(N, size, rank)
n, s = blockdist(N, size, rank)
index[max_axis] = np.s_[s:s+n]
index = tuple(index)
data = data[index]
Expand All @@ -119,7 +119,7 @@ def distribute_data(data, comm):
shape = data.shape
index = [np.s_[:]] * len(shape)
N = shape[max_axis+1]
n, s = _blockdist(N, size, rank)
n, s = blockdist(N, size, rank)
index[max_axis+1] = np.s_[s:s+n]
index = tuple(index)
data = data[index]
Expand All @@ -141,7 +141,7 @@ def distribute_dimension(data, max_axis, comm):
shape = data.shape
index = [np.s_[:]] * len(shape)
N = shape[max_axis]
n, s = _blockdist(N, size, rank)
n, s = blockdist(N, size, rank)
index[max_axis] = np.s_[s:s+n]
index = tuple(index)
data = data[index]
Expand All @@ -151,7 +151,7 @@ def distribute_dimension(data, max_axis, comm):
return data


def _blockdist(N, size, rank):
def blockdist(N, size, rank):
q, r = divmod(N, size)
n = q + (1 if r > rank else 0)
s = rank * q + min(rank, r)
Expand Down
89 changes: 60 additions & 29 deletions pyspod/utils/reader.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -162,6 +162,15 @@ def is_real(self):
# each process has complete data for a subset of space.
#
########################################################################################
def _inspect_file_metadata(filename, variable):
with xr.open_dataset(filename, cache=False) as dataset:
data = dataset[variable]
shape = data.shape
dtype = data.dtype

size_gb = os.path.getsize(filename) / 1024 / 1024 / 1024
return shape, dtype, size_gb

class reader_2stage():
def __init__(self, data_list, xdim, dtype, comm, nv, variables, nreaders = None, nchunks = 3, nblocks = 3):
assert comm is not None, "2-stage reader requires MPI"
Expand Down Expand Up @@ -195,20 +204,42 @@ def __init__(self, data_list, xdim, dtype, comm, nv, variables, nreaders = None,
self._max_axes = np.array([1,0]) # time is the first dimension (not listed), then nvar, then the spatial dimension
nt = 0
shape = None

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):
f = data_list[i]
file_shape, file_dtype, file_size = _inspect_file_metadata(f, variables[0])
local_infos.append((f, file_shape, str(file_dtype), file_size))

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

if comm.rank == 0:
for f in data_list:
d = xr.open_dataset(f,cache=False)[variables[0]]
# make sure that all files have the same spatial shape
infos = []
for rank_infos in all_infos:
infos.extend(rank_infos)

nt = 0
shape = None

for f, file_shape, file_dtype, file_size in infos:
# Make sure that all files have the same spatial shape
if shape is not None:
assert d.shape[1:] == shape[1:], f'File {f} has different shape than the previous ones'
shape = d.shape
self._file_time[f] = (nt, nt+shape[0])
nt += shape[0]
if d.dtype != 'float32' and d.dtype != 'float64':
self._is_real = False
d.close()
self._files_size += os.path.getsize(f)/1024/1024/1024 # GB
assert file_shape[1:] == shape[1:], f'File {f} has different shape than the previous ones'

shape = file_shape

self._file_time[f] = (nt, nt + file_shape[0])
nt += file_shape[0]

if file_dtype != 'float32' and file_dtype != 'float64':
self._is_real = False

self._files_size += file_size

self._shape = (nt,) + shape[1:] + (self._nv,)

self._shape = comm.bcast(self._shape, root=0)
Expand Down Expand Up @@ -244,12 +275,12 @@ def get_data_for_time(self, ts, te):
n_all_xyz = np.prod(self._shape[1:-1]) # product of spatial dimensions

# first distribute by the time dimension to maximize contiguous reads and minimize the number of readers per file
n_dist_time, s_dist_time = utils_par._blockdist(te-ts, self._nreaders, mpi_rank)
n_dist_time, s_dist_time = utils_par.blockdist(te-ts, self._nreaders, mpi_rank)
js = ts + s_dist_time
je = ts + s_dist_time+n_dist_time

# allocate local data
n_dist_xyz, _ = utils_par._blockdist(n_all_xyz, mpi_size, mpi_rank)
n_dist_xyz, _ = utils_par.blockdist(n_all_xyz, mpi_size, mpi_rank)
self._local_shape = n_dist_xyz

cum_t = 0
Expand Down Expand Up @@ -290,7 +321,7 @@ def get_data_for_time(self, ts, te):
# each process receives n0: every other process's len(time)
# x n1: own slice of n_all_xyz OR process 0-sized slice of n_all_xyz (when padding is used)

max_dist_xyz, _ = utils_par._blockdist(n_all_xyz, mpi_size, 0) # proc 0 has the largest number of elements
max_dist_xyz, _ = utils_par.blockdist(n_all_xyz, mpi_size, 0) # proc 0 has the largest number of elements
max_dist_time = comm.allreduce(n_dist_time, op=MPI.MAX)

use_padding = False if MPI.VERSION >= 4 or max_dist_time*mpi_size*max_dist_xyz*self._nv < np.iinfo(np.int32).max else True
Expand All @@ -309,10 +340,10 @@ def get_data_for_time(self, ts, te):

offset = np.zeros(mpi_size+1, dtype=np.int32)
for irank in range(mpi_size):
nt, _ = utils_par._blockdist(te-ts, self._nreaders, irank) # time (distributed on the reader)
nt, _ = utils_par.blockdist(te-ts, self._nreaders, irank) # time (distributed on the reader)
offset[irank+1] = offset[irank] + nt

n_irank, s_irank = utils_par._blockdist(n_all_xyz, mpi_size, irank)
n_irank, s_irank = utils_par.blockdist(n_all_xyz, mpi_size, irank)
s_msgs[irank] = input_data[:,s_irank:s_irank+n_irank,:].copy()

del input_data
Expand Down Expand Up @@ -340,10 +371,10 @@ def get_data_for_time(self, ts, te):
s_msgs = {}

for irank in range(mpi_size):
nt, _ = utils_par._blockdist(te-ts, self._nreaders, irank) # time (distributed on the reader)
nt, _ = utils_par.blockdist(te-ts, self._nreaders, irank) # time (distributed on the reader)
recvcounts[irank] = nt*self._nv

irank_n_xyz, irank_s_xyz = utils_par._blockdist(n_all_xyz, mpi_size, irank)
irank_n_xyz, irank_s_xyz = utils_par.blockdist(n_all_xyz, mpi_size, irank)
s_msgs[irank] = np.zeros((input_data.shape[0],max_dist_xyz,input_data.shape[2]),dtype=self._dtype)
s_msgs[irank][:,:irank_n_xyz,:] = (input_data[:,irank_s_xyz:irank_s_xyz+irank_n_xyz,:]) # nmax-sized and 0-padded (if needed) array
del input_data
Expand Down Expand Up @@ -383,15 +414,15 @@ def get_data(self, ts = None):

for chunk in range(0,nchunks):

t_n, t_s = utils_par._blockdist(self.nt, nchunks, chunk)
t_n, t_s = utils_par.blockdist(self.nt, nchunks, chunk)
t_e = t_s + t_n

x = self.get_data_for_time(t_s,t_e)

for blk in range(0,nblks):
blk_idx = chunk*nblks + blk

blk_t_n, blk_t_s = utils_par._blockdist(t_e-t_s, nblks, blk)
blk_t_n, blk_t_s = utils_par.blockdist(t_e-t_s, nblks, blk)
blk_t_e = blk_t_s + blk_t_n

data_dict[blk_idx] = {}
Expand Down Expand Up @@ -541,12 +572,12 @@ def get_data_for_time(self, ts, te):
n_all_xyz = np.prod(self._shape[1:-1]) # product of spatial dimensions

# fist distribute by the time dimension to maximize contiguous reads and minimize the number of readers per file
n_dist_time, s_dist_time = utils_par._blockdist(te-ts, mpi_size, mpi_rank)
n_dist_time, s_dist_time = utils_par.blockdist(te-ts, mpi_size, mpi_rank)
js = ts + s_dist_time
je = ts + s_dist_time+n_dist_time

# allocate local data
n_dist_xyz, _ = utils_par._blockdist(n_all_xyz, mpi_size, mpi_rank)
n_dist_xyz, _ = utils_par.blockdist(n_all_xyz, mpi_size, mpi_rank)
self._local_shape = n_dist_xyz

cum_t = 0
Expand Down Expand Up @@ -583,7 +614,7 @@ def get_data_for_time(self, ts, te):
# each process receives n0: every other process's len(time)
# x n1: own slice of n_all_xyz OR process 0-sized slice of n_all_xyz (when padding is used)

max_dist_xyz, _ = utils_par._blockdist(n_all_xyz, mpi_size, 0) # proc 0 has the largest number of elements
max_dist_xyz, _ = utils_par.blockdist(n_all_xyz, mpi_size, 0) # proc 0 has the largest number of elements
max_dist_time = comm.allreduce(n_dist_time, op=MPI.MAX)

use_padding = False if MPI.VERSION >= 4 or max_dist_time*mpi_size*max_dist_xyz*self._nv < np.iinfo(np.int32).max else True
Expand All @@ -602,10 +633,10 @@ def get_data_for_time(self, ts, te):

offset = np.zeros(mpi_size+1, dtype=np.int32)
for irank in range(mpi_size):
nt, _ = utils_par._blockdist(te-ts, mpi_size, irank) # time (distributed on the reader)
nt, _ = utils_par.blockdist(te-ts, mpi_size, irank) # time (distributed on the reader)
offset[irank+1] = offset[irank] + nt

n_irank, s_irank = utils_par._blockdist(n_all_xyz, mpi_size, irank)
n_irank, s_irank = utils_par.blockdist(n_all_xyz, mpi_size, irank)
s_msgs[irank] = input_data[:,s_irank:s_irank+n_irank,:].copy()

del input_data
Expand Down Expand Up @@ -633,10 +664,10 @@ def get_data_for_time(self, ts, te):
s_msgs = {}

for irank in range(mpi_size):
nt, _ = utils_par._blockdist(te-ts, mpi_size, irank) # time (distributed on the reader)
nt, _ = utils_par.blockdist(te-ts, mpi_size, irank) # time (distributed on the reader)
recvcounts[irank] = nt*self._nv#*max_dist_xyz

irank_n_xyz, irank_s_xyz = utils_par._blockdist(n_all_xyz, mpi_size, irank)
irank_n_xyz, irank_s_xyz = utils_par.blockdist(n_all_xyz, mpi_size, irank)
s_msgs[irank] = np.zeros((input_data.shape[0],max_dist_xyz,input_data.shape[2]),dtype=self._dtype)
s_msgs[irank][:,:irank_n_xyz,:] = (input_data[:,irank_s_xyz:irank_s_xyz+irank_n_xyz,:]) # nmax-sized and 0-padded (if needed) array
del input_data
Expand Down Expand Up @@ -676,15 +707,15 @@ def get_data(self, ts = None):

for chunk in range(0,nchunks):

t_n, t_s = utils_par._blockdist(self.nt, nchunks, chunk)
t_n, t_s = utils_par.blockdist(self.nt, nchunks, chunk)
t_e = t_s + t_n

x = self.get_data_for_time(t_s,t_e)

for blk in range(0,nblks):
blk_idx = chunk*nblks + blk

blk_t_n, blk_t_s = utils_par._blockdist(t_e-t_s, nblks, blk)
blk_t_n, blk_t_s = utils_par.blockdist(t_e-t_s, nblks, blk)
blk_t_e = blk_t_s + blk_t_n

data_dict[blk_idx] = {}
Expand Down
Loading