def to_hdf5(self, f, dataset, create=False):
"""Parallel write into a contiguous HDF5 dataset.
Parameters
----------
filename : str, h5py.File or h5py.Group
File to write dataset into.
dataset : string
Name of dataset to write into. Should not exist.
"""
## Naive non-parallel implementation to start
import h5py
if not h5py.get_config().mpi:
if isinstance(f, basestring):
self._to_hdf5_serial(f, dataset, create)
return
else:
raise ValueError(
"Argument must be a filename if h5py does not have MPI support"
)
mode = 'a' if create else 'r+'
fh = misc.open_h5py_mpi(f, mode, self.comm)
dset = fh.create_dataset(dataset, shape=self.global_shape, dtype=self.dtype)
start = self.local_offset[self.axis]
end = start + self.local_shape[self.axis]
# Construct slices for axis
sl = [slice(None, None)] * self.axis
sl += [slice(start, end)]
sl = tuple(sl)
# Check that there are no null slices, otherwise we need to turn off
# collective IO to work around an h5py issue (#965)
no_null_slices = self.global_shape[self.axis] >= self.comm.size
if fh.is_mpi and no_null_slices:
with dset.collective:
dset[sl] = self[:]
else:
dset[sl] = self[:]
if fh.opened:
fh.close()
def to_hdf5(
self,
f,
dataset,
create=False,
chunks=None,
compression=None,
compression_opts=None,
):
"""Parallel write into a contiguous HDF5 dataset.
Parameters
----------
filename : str, h5py.File or h5py.Group
File to write dataset into.
dataset : string
Name of dataset to write into. Should not exist.
"""
import h5py
if not h5py.get_config().mpi:
if isinstance(f, basestring):
self._to_hdf5_serial(f, dataset, create)
return
else:
raise ValueError(
"Argument must be a filename if h5py does not have MPI support"
)
mode = "a" if create else "r+"
fh = misc.open_h5py_mpi(f, mode, self.comm)
start = self.local_offset[self.axis]
end = start + self.local_shape[self.axis]
# Construct slices for axis
sel = ([slice(None, None)] * self.axis) + [slice(start, end)]
sel = _expand_sel(sel, self.ndim)
# Check that there are no null slices, otherwise we need to turn off
# collective IO to work around an h5py issue (#965)
no_null_slices = self.global_shape[self.axis] >= self.comm.size
# Split the axis to get the IO size under ~2GB (only if MPI-IO)
split_axis, partitions = self._partition_io(skip=(not fh.is_mpi))
# Only use collective IO if:
# - there are no null slices (h5py bug)
# - we are not distributed over axis=0 as there is no advantage for
# collective IO which is usually slow
# - unless we want to use compression/chunking
# TODO: change if h5py bug fixed
# TODO: better would be a test on contiguous IO size
use_collective = (fh.is_mpi and no_null_slices
and (self.axis > 0 or compression is not None))
if fh.is_mpi and not use_collective:
# Need to disable compression if we can't use collective IO
chunks, compression, compression_opts = None, None, None
dset = fh.create_dataset(
dataset,
shape=self.global_shape,
dtype=self.dtype,
chunks=chunks,
compression=compression,
compression_opts=compression_opts,
)
# Read using collective MPI-IO if specified
with dset.collective if use_collective else DummyContext():
# Loop over partitions of the IO and perform them
for part in partitions:
islice, fslice = _partition_sel(sel, split_axis,
self.global_shape[split_axis],
part)
dset[islice] = self[fslice]
if fh.opened:
fh.close()
def from_hdf5(cls, f, dataset, comm=None, axis=0, sel=None):
"""Read MPIArray from an HDF5 dataset in parallel.
Parameters
----------
f : filename, or `h5py.File` object
File to read dataset from.
dataset : string
Name of dataset to read from. Must exist.
comm : MPI.Comm, optional
MPI communicator to distribute over. If `None` optional, use
`MPI.COMM_WORLD`.
axis : int, optional
Axis over which the read should be distributed. This can be used
to select the most efficient axis for the reading.
sel : tuple, optional
A tuple of slice objects used to make a selection from the array
*before* reading. The output will be this selection from the dataset
distributed over the given axis.
Returns
-------
array : MPIArray
"""
fh = misc.open_h5py_mpi(f, 'r', comm)
dset = fh[dataset]
dshape = dset.shape # Shape of the underlying dataset
naxis = len(dshape)
dtype = dset.dtype
# Check that the axis is valid and wrap to an actual position
if axis < -naxis or axis >= naxis:
raise ValueError("Distributed axis %i not in range (%i, %i)" % (axis, -naxis, naxis-1))
axis = naxis + axis if axis < 0 else axis
# Ensure sel is defined to cover all axes
if sel is None:
sel = [slice(None)] * naxis
if len(sel) < naxis:
sel = list(sel) + [slice(None)] * (naxis - len(sel))
sel = list(sel)
# Figure out the final array size and create it
gshape = []
for l, sl in zip(dshape, sel):
start, end, stride = sl.indices(l)
n = 1 + (end - start - 1) // stride
gshape.append(n)
dist_arr = cls(gshape, axis=axis, comm=comm, dtype=dtype)
# Get the local start and end indices
lstart = dist_arr.local_offset[axis]
lend = lstart + dist_arr.local_shape[axis]
# Create the slice object into the dataset by resolving the rank's slice on the sel
dstart, dend, dstride = sel[axis].indices(dshape[axis])
sel[axis] = slice(dstart + lstart * dstride, dstart + lend * dstride, dstride)
sel = tuple(sel)
# Check that there are no null slices, otherwise we need to turn off
# collective IO to work around an h5py issue (#965)
no_null_slices = dist_arr.global_shape[axis] >= dist_arr.comm.size
# Read using MPI-IO if possible
if fh.is_mpi and no_null_slices:
with dset.collective:
dist_arr[:] = dset[sel]
else:
dist_arr[:] = dset[sel]
if fh.opened:
fh.close()
return dist_arr
def from_hdf5(cls, f, dataset, comm=None, axis=0, sel=None):
"""Read MPIArray from an HDF5 dataset in parallel.
Parameters
----------
f : filename, or `h5py.File` object
File to read dataset from.
dataset : string
Name of dataset to read from. Must exist.
comm : MPI.Comm, optional
MPI communicator to distribute over. If `None` optional, use
`MPI.COMM_WORLD`.
axis : int, optional
Axis over which the read should be distributed. This can be used
to select the most efficient axis for the reading.
sel : tuple, optional
A tuple of slice objects used to make a selection from the array
*before* reading. The output will be this selection from the dataset
distributed over the given axis.
Returns
-------
array : MPIArray
"""
from mpi4py import MPI
# Don't both using MPI where the axis is not zero. It's probably just slower.
# TODO: with tuning this might not be true. Keep an eye on this.
use_mpi = axis > 0
# Read the file. Opening with MPI if requested, and we can
fh = misc.open_h5py_mpi(f, "r", use_mpi=use_mpi, comm=comm)
dset = fh[dataset]
dshape = dset.shape # Shape of the underlying dataset
naxis = len(dshape)
dtype = dset.dtype
# Check that the axis is valid and wrap to an actual position
if axis < -naxis or axis >= naxis:
raise ValueError("Distributed axis %i not in range (%i, %i)" %
(axis, -naxis, naxis - 1))
axis = naxis + axis if axis < 0 else axis
# Ensure sel is defined to cover all axes
sel = _expand_sel(sel, naxis)
# Figure out the final array size and create it
gshape = []
for l, sl in zip(dshape, sel):
gshape.append(_len_slice(sl, l))
dist_arr = cls(gshape, axis=axis, comm=comm, dtype=dtype)
# Get the local start and end indices
lstart = dist_arr.local_offset[axis]
lend = lstart + dist_arr.local_shape[axis]
# Create the slice object into the dataset by resolving the rank's slice on the sel
sel[axis] = _reslice(sel[axis], dshape[axis], slice(lstart, lend))
sel = tuple(sel)
# Split the axis to get the IO size under ~2GB (only if MPI-IO)
split_axis, partitions = dist_arr._partition_io(skip=(not fh.is_mpi))
# Check that there are no null slices, otherwise we need to turn off
# collective IO to work around an h5py issue (#965)
no_null_slices = dist_arr.global_shape[axis] >= dist_arr.comm.size
# Only use collective IO if:
# - there are no null slices (h5py bug)
# - we are not distributed over axis=0 as there is no advantage for
# collective IO which is usually slow
# TODO: change if h5py bug fixed
# TODO: better would be a test on contiguous IO size
# TODO: do we need collective IO to read chunked data?
use_collective = fh.is_mpi and no_null_slices and axis > 0
# Read using collective MPI-IO if specified
with dset.collective if use_collective else DummyContext():
# Loop over partitions of the IO and perform them
for part in partitions:
islice, fslice = _partition_sel(sel, split_axis,
gshape[split_axis], part)
dist_arr[fslice] = dset[islice]
if fh.opened:
fh.close()
return dist_arr