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