mpi4py 进阶之 memh5

在上一篇中我们介绍了利用 mpi4py 和 h5py 进行并行分布式的 HDF5 文件操作,下面我们将介绍 caput 中的 memh5 模块,其中提供若干功能强大的工具可以使我们在内存中操作与 HDF5 文件类似的数据结构,并提供该数据结构到磁盘中 HDF5 文件之间的映射,当然这些操作都是可以并行分布式地进行的。

类及操作方法

能够一致地操作无论是内存中还是磁盘上文件中的数据在某些情况下是非常有用的,而且也能为我们的编程带来很大的方便。我们前面介绍的 HDF5 文件是一种非常强大的科学数据文件存储格式,其 Python 操作工具 h5py 使我们能够非常方便地使用 HDF5 文件及操作其中存放的数据,结合 mpi4py 还能进行并行分布式的 HDF5 文件操作。但是作为一种文件格式,HDF5 文件及其中存放的数据是存在于磁盘中的,用 h5py 对其进行操作时,其中存放的数据(包括 dataset 和 attribute)一般读取到内存后会成为一个个独立的数据项,比如说 numpy 数组,数字,字符串等,从而失去其在原文件中的组织形式和等级结构。内存中的这些数据项的操作方式和 h5py 提供的对 HDF5 文件的操作方式并不一致,这给我们带来很多不便。如果能够将 HDF5 文件中的数据保持其在原文件中的组织形式和等级结构复制到内存中,建立一个内存中的数据到原文件中的数据的映射,并提供一套与 HDF5 文件操作一致的方法来操作内存中的数据映像,必然给我们的编程带来莫大的方便,也使得 HDF5 文件及其数据更加易用。幸运的是,caput 中的 memh5 模块给我们提供了这些工具,下面给出其提供的主要功能及其方法接口。

MemGroup

内存中的 group 实现,对应于 h5py 中的 Group。

class MemGroup(distributed=False, comm=None)

创建一个 MemGroup 对象,如果 distributed 为 True,且 comm 为一个有效的通信子对象,则可以在其中创建分布式的 MemDatasetDistributed (见下面的介绍)。在 mpi4py 可用的情况下,如果 comm = None, 则会使用 MPI.COMM_WORLD。

除了与 h5py.Group 一致的方法外,MemGroup 提供以下额外的或不同的方法接口:

from_group(cls, group)

深复制 group 而创建一个 MemGroup。group 可以为一个 MemGroup 或者 h5py.Group (包括 h5py.File,也即其 "/" group) 对象。

from_hdf5(cls, filename, distributed=False, hints=True, comm=None, **kwargs)

读取整个 HDF5 文件来创建一个 MemGroup。filename 为 HDF5 文件名,distributed 指明是否创建一个分布式的 MemGroup,即其中可以包含一些 MemDatasetDistributed,如果 distributed 为 True,hints 为 True,且 comm 是一个有效的通信子对象,则会根据文件中 dataset 的 '__memh5_distributed_dset' 属性(如果存在且为 True)将该 dataset 读取到内存中创建为一个 MemDatasetDistributed,其它的 dataset 则读取到内存后被创建为 MemDatasetCommon。

to_hdf5(self, filename, hints=True, **kwargs)

将该 MemGroup 中的数据按照其组织形式及等级结构写入到一个 HDF5 文件 filename 中,hints 以指明是否在文件中标注哪些 dataset 是由该 MemGroup 中的 MemDatasetDistributed 所存储,以变于 from_hdf5 方法可将其恢复为 MemDatasetDistributed。

create_dataset(self, name, shape=None, dtype=None, data=None, distributed=False, distributed_axis=None, **kwargs)

在该 MemGroup 中创建一个 MemDataset。nameshapedtype 同 h5py 对应函数的对应参数,如果 distributed 为 True,则会创建一个分布在轴 distributed_axis (0 如果为 None) 上的 MemDatasetDistributed,此时如果 data 非 None,则其必须是一个 MPIArray;如果 distributed 为 False,则会创建一个 MemDatasetCommon。

dataset_common_to_distributed(self, name, distributed_axis=0)

将该 MemGroup 中命名为 name 的 MemDatasetCommon 转化为分布在轴 distributed_axis 上的 MemDatasetDistributed。如果此 dataset 本身就是 MemDatasetDistributed,则会将其重新分布到 distributed_axis 轴上。

dataset_distributed_to_common(self, name)

将该 MemGroup 中命名为 name 的 MemDatasetDistributed 转化为 MemDatasetCommon。

MemDataset

内存中的 dataset 实现,对应于 h5py 中的 Dataset。

class MemDataset(**kwargs)

它是一个抽象基类,不能直接使用,应该使用其子类 MemDatasetCommonMemDatasetDistributed

MemDatasetCommon

class MemDatasetCommon(shape, dtype)

shapedtype 参数创建一个全 0 的 MemDatasetCommon。MemDatasetCommon 是指在每个进程中都存在且完全一样的 MemDataset (其数据为一个 numpy array)。一般对非常大的 numpy array,不适宜将其创建为 MemDatasetCommon,因为会占用大量的内存,可用将其创建为 MemDatasetDistributed。

MemDatasetCommon 除了与 h5py.Dataset 接口一致的方法外,还有下面这个方法:

from_numpy_array(cls, data)

由一个 numpy array data 创建一个 MemDatasetCommon。

MemDatasetCommon 除了与 h5py.Dataset 一致的属性外,还有以下属性:

comm

为 None。

common

为 True。

distributed

为 False。

data

其包含的数据,为一个 numpy array。

local_data

同 data。

MemDatasetDistributed

class MemDatasetDistributed(shape, dtype, axis=0, comm=None)

由参数 sahpedtype,分布轴 axis 和通信子对象 comm 创建一个空的 MemDatasetDistributed。MemDatasetDistributed 是指沿某个轴分布在各个进程上的 MemDataset (其数据为一个 MPIArray),每个进程持有该 MPIArray 的一个子数组。

MemDatasetDistributed 除了与 h5py.Dataset 接口一致的方法外,还有下面这些方法:

from_mpi_array(cls, data)

由一个 MPIArray data 创建一个 MemDatasetDistributed。

redistribute(self, axis)

改变该 MemDatasetDistributed 的分布轴。

MemDatasetDistributed 除了与 h5py.Dataset 一致的属性外,还有以下属性:

comm

为该 MemDatasetDistributed 所分布的通信子。

common

为 False。

distributed

为 True。

distributed_axis

分布轴。

data

其包含的数据,为一个 MPIArray。

local_data

每个进程所持有的子数组,为一个 numpy array。

global_shape

其包含的 MPIArray 的 global_shape。

local_shape

每个进程所持有的子数组的 shape。

local_offset

每个进程所持有的子数组在其 MPIArray 中各维的偏移量。

MemAttrs

内存中的 attribute 实现,对应于 h5py 中的 AttributeManager。

class MemAttrs(**kwargs):

创建一个 MemAttrs,实际上为一个 Python 字典。

MemDiskGroup

一个更高级的数据容器,其数据可能存在于磁盘上(一个 HDF5 文件),也可能在内存中(一个 MemGroup)。

class MemDiskGroup(data_group=None, distributed=False, comm=None)

创建一个 MemDiskGroup。data_group 可以为 None,字符串,h5py.Group 或 MemGroup,如果为 None,会先创建一个空的 MemGroup 并以此来创建 MemDiskGroup,如果为一个字符串,则会打开一个以该字符串命名的 HDF5 文件并以该文件句柄来创建一个 MemDiskGroup,如果是 MemGroup,则参数 distributedcomm 指明是否能在其中创建或容纳 MemDatasetDistributed。

MemDiskGroup 本身可能并不是特别有用,其主要目的是作为构建更复杂更具体数据容器的基类,因此不做进一步介绍。

若干工具函数

memh5 模块提供的若干工具函数如下:

attrs2dict(attrs)

将 h5py attributes attrs 深复制到一个字典中,返回此字典。

is_group(obj)

测试 obj 是否是 group (h5py.Group 或 MemGroup)。

copyattrs(a1, a2)

a1 深复制到 a2 中。a1a2 都可以为 h5py.AttributeManager 或者 MemAttrs。

deep_group_copy(g1, g2)

g1 深复制到 g2 中。g1g2 都可以为 h5py.Group 或 MemGroup。

例程

下面给出简单的使用例程。

# memh5_demo.py

"""
Demonstrates how to use memh5.

Run this with 4 processes like:
$ mpiexec -n 4 python memh5_demo.py
"""

import os
import numpy as np
import h5py
from mpi4py import MPI
from caput import memh5


comm = MPI.COMM_WORLD
rank = comm.rank
size = comm.size

N = 10
data = np.arange(N, dtype=np.int32)

file_name = 'test.hdf5'
# create a new HDF5 file collectively
f = h5py.File(file_name, driver='mpio', comm=comm)
# use atomic mode
f.atomic = True
# create a new group collectively
f.create_group('grp')
# create a dataset in group "/grp" collectively
f.create_dataset('grp/dset', data=data)
# set an attribute of the dataset grp/dset collectively
f['grp/dset'].attrs['a'] = 1
# close file collectively
f.close()

# create MemGroup from the HDF5 file
mgrp = memh5.MemGroup.from_hdf5(file_name, distributed=True, comm=comm)
# create a new MemDatasetCommon
dset1 = mgrp.create_dataset('dset1', shape=(3, 4), dtype=np.float32)
print 'dset1.common = %s, dset1.distributed = %s' % (dset1.common, dset1.distributed)
# create a new MemDatasetDistributed
dset2 = mgrp.create_dataset('dset2', shape=(3, 4), dtype=np.float32, distributed=True, distributed_axis=1)
print 'dset2.common = %s, dset2.distributed = %s, dset2.distributed_axis = %d' % (dset2.common, dset2.distributed, dset2.distributed_axis)
# redistribute dset2 to axis 0
dset2.redistribute(axis=0)
print 'after redistribute, dset2.distributed_axis = %d' %  dset2.distributed_axis
# change dset2 to be MemDatasetCommon
mgrp.dataset_distributed_to_common('dset2')
# create and set an attribute of dset2
dset2.attrs['status'] = 'new crated'
# write mgrp to a HDF5 file
mgrp.to_hdf5('new.hdf5')

# remove the created files
if rank == 0:
    os.remove(file_name)
    os.remove('new.hdf5')

运行结果如下:

$ mpiexec -n 4 python memh5_demo.py
Starting MPI rank=0 [size=4]
Starting MPI rank=2 [size=4]
Starting MPI rank=3 [size=4]
Starting MPI rank=1 [size=4]
dset1.common = True, dset1.distributed = False
dset2.common = False, dset2.distributed = True, dset2.distributed_axis = 1
dset1.common = True, dset1.distributed = False
dset2.common = False, dset2.distributed = True, dset2.distributed_axis = 1
dset1.common = True, dset1.distributed = False
dset2.common = False, dset2.distributed = True, dset2.distributed_axis = 1
dset1.common = True, dset1.distributed = False
dset2.common = False, dset2.distributed = True, dset2.distributed_axis = 1
after redistribute, dset2.distributed_axis = 0
after redistribute, dset2.distributed_axis = 0
after redistribute, dset2.distributed_axis = 0
after redistribute, dset2.distributed_axis = 0

以上我们介绍了 caput 中的 memh5 模块,在下一篇中我们将介绍一个建立在 mpi4py 基础之上的有用工具 mpipool,允许我们非常方便地将一个函数并行地应用到一系列数据上。

你可能感兴趣的:(mpi4py 进阶之 memh5)