From 3ff6ccb37326bb3e895d3f154a3e547dd328e7e9 Mon Sep 17 00:00:00 2001 From: baarn Date: Thu, 27 Aug 2026 10:51:58 -0400 Subject: [PATCH 1/2] Parallelize file metadata inspection in 2-stage reader --- pyspod/utils/reader.py | 59 ++++++++++++++++++++++++++++++++++-------- 1 file changed, 48 insertions(+), 11 deletions(-) diff --git a/pyspod/utils/reader.py b/pyspod/utils/reader.py index 0138429..fe85621 100644 --- a/pyspod/utils/reader.py +++ b/pyspod/utils/reader.py @@ -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" @@ -195,20 +204,48 @@ 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 + + # Number of ranks used to inspect file metadata + metadata_nreaders = min(self._nreaders, len(data_list)) + + local_infos = [] + + # Distribute files across reader ranks + if comm.rank < metadata_nreaders: + for i in range(comm.rank, len(data_list), metadata_nreaders): + f = data_list[i] + file_shape, file_dtype, file_size = _inspect_file_metadata(f, variables[0]) + local_infos.append((i, 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) + + # Restore the original file order + infos.sort(key=lambda x: x[0]) + + 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) From 1ee2f69ebbc13060fc675f46d9a5def3d091a2e6 Mon Sep 17 00:00:00 2001 From: baarn Date: Sun, 30 Aug 2026 08:32:33 -0400 Subject: [PATCH 2/2] Use block distribution for metadata inspection --- pyspod/utils/parallel.py | 10 +++---- pyspod/utils/reader.py | 62 ++++++++++++++++++---------------------- 2 files changed, 33 insertions(+), 39 deletions(-) diff --git a/pyspod/utils/parallel.py b/pyspod/utils/parallel.py index 229686c..a9584ee 100644 --- a/pyspod/utils/parallel.py +++ b/pyspod/utils/parallel.py @@ -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() @@ -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] @@ -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] @@ -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] @@ -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) diff --git a/pyspod/utils/reader.py b/pyspod/utils/reader.py index fe85621..96d42d6 100644 --- a/pyspod/utils/reader.py +++ b/pyspod/utils/reader.py @@ -204,19 +204,16 @@ 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 - - # Number of ranks used to inspect file metadata - metadata_nreaders = min(self._nreaders, len(data_list)) - + local_infos = [] - + # Distribute files across reader ranks - if comm.rank < metadata_nreaders: - for i in range(comm.rank, len(data_list), metadata_nreaders): - f = data_list[i] - file_shape, file_dtype, file_size = _inspect_file_metadata(f, variables[0]) - local_infos.append((i, f, file_shape, str(file_dtype), file_size)) - + 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) @@ -224,14 +221,11 @@ def __init__(self, data_list, xdim, dtype, comm, nv, variables, nreaders = None, infos = [] for rank_infos in all_infos: infos.extend(rank_infos) - - # Restore the original file order - infos.sort(key=lambda x: x[0]) - + nt = 0 shape = None - for _, f, file_shape, file_dtype, file_size in infos: + 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 file_shape[1:] == shape[1:], f'File {f} has different shape than the previous ones' @@ -281,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 @@ -327,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 @@ -346,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 @@ -377,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 @@ -420,7 +414,7 @@ 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) @@ -428,7 +422,7 @@ def get_data(self, ts = None): 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] = {} @@ -578,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 @@ -620,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 @@ -639,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 @@ -670,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 @@ -713,7 +707,7 @@ 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) @@ -721,7 +715,7 @@ def get_data(self, ts = None): 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] = {}