bcast_test.py

import hashlib
import sys

import numpy as np
from mpi4py import MPI

comm = MPI.COMM_WORLD
UINT32_MAX = 4294967295
INT32_MAX = 2147483647

REPEAT = 10
ROOT = 0

def calc_checksum(array):
    h = hashlib.md5()
    h.update(array)
    return h.hexdigest()


def init_buf(num_elem):
    if comm.rank == 0:
        return np.random.randint(0, INT32_MAX + 1, num_elem, dtype=np.int32)
    else:
        return np.zeros(num_elem, dtype=np.int32)


def check_result(checksum, iter):
    if comm.rank == ROOT:
        peer_checksum = comm.recv(source=1 - ROOT)
        if peer_checksum != checksum:
            print("Error: checksum mismatch. ")
            print("       rank {}: {}".format(ROOT, checksum))
            print("       rank {}: {}".format(1 - ROOT, peer_checksum))
        else:
            print("Iteration {} OK.".format(iter))
    else:
        comm.send(checksum, dest=ROOT)


def main():
    assert comm.size == 2

    if len(sys.argv) > 1:
        buf_size = int(sys.argv[1])
    else:
        buf_size = 32 * 1024 * 1024

    num_elem = buf_size // 4  # int32

    for i in range(REPEAT):
        array = init_buf(num_elem)
        comm.Bcast(array, root=ROOT)
        checksum = calc_checksum(array)
        check_result(checksum, i)


if __name__ == '__main__':
    main()