1# This file is to test collective io in h5py
2
3"""
4Author:  Jialin Liu, jalnliu@lbl.gov
5Date:    Nov 17, 2015
6Prerequisites: python 2.5.0, mpi4py and numpy
7Source Codes: Already submit this 'collective io' branch to h5py master, meanwhile, can download this branch at https://github.com/valiantljk/h5py.git
8Note: Must build the h5py with parallel hdf5
9"""
10
11from mpi4py import MPI
12import numpy as np
13import h5py
14import sys
15
16#"run as "mpirun -np 64 python-mpi collective_io.py 1 file.h5"
17#(1 is for collective write, other numbers for non-collective write)"
18
19colw=1 #default is collective write
20filename="parallel_test.hdf5"
21if len(sys.argv)>2:
22    colw = int(sys.argv[1])
23    filename=str(sys.argv[2])
24comm =MPI.COMM_WORLD
25nproc = comm.Get_size()
26f = h5py.File(filename, 'w', driver='mpio', comm=MPI.COMM_WORLD)
27rank = comm.Get_rank()
28length_x = 6400*1024
29length_y = 1024
30dset = f.create_dataset('test', (length_x,length_y), dtype='f8')
31#data type should be consistent in numpy and h5py, e.g., 64 bits
32#otherwise, hdf5 layer will fall back to independent io.
33f.atomic = False
34length_rank=length_x / nproc
35length_last_rank=length_x -length_rank*(nproc-1)
36comm.Barrier()
37timestart=MPI.Wtime()
38start=rank*length_rank
39end=start+length_rank
40if rank==nproc-1: #last rank
41    end=start+length_last_rank
42temp=np.random.random((end-start,length_y))
43if colw==1:
44    with dset.collective:
45        dset[start:end,:] = temp
46else:
47    dset[start:end,:] = temp
48comm.Barrier()
49timeend=MPI.Wtime()
50if rank==0:
51    if colw==1:
52        print("collective write time %f" %(timeend-timestart))
53    else:
54        print("independent write time %f" %(timeend-timestart))
55    print("data size x: %d y: %d" %(length_x, length_y))
56    print("file size ~%d GB" % (length_x*length_y/1024.0/1024.0/1024.0*8.0))
57    print("number of processes %d" %nproc)
58f.close()
59