def distribute_list(source):
"""Distributes the list from root to individual nodes
"""
# quick check
if SIZE == 1:
return source
if is_root():
length = len(source)
if length == 0:
logging.warning("Warning: List has length 0")
else:
length = 0
length = COMM.bcast(length)
if length == 0:
return []
segments = get_segments(length)
if is_root():
for i in range(1,SIZE):
send_list = source[segments[i]:segments[i+1]]
COMM.send(send_list, dest=i)
data = source[:segments[1]]
del source
else:
data = COMM.recv()
return data
def distribute_list(source):
"""Distributes the list from root to individual nodes
"""
# quick check
if SIZE == 1:
return source
if is_root():
length = len(source)
if length == 0:
logging.warning("Warning: List has length 0")
else:
length = 0
length = COMM.bcast(length)
if length == 0:
return []
segments = get_segments(length)
if is_root():
for i in range(1, SIZE):
send_list = source[segments[i]:segments[i + 1]]
COMM.send(send_list, dest=i)
data = source[:segments[1]]
del source
else:
data = COMM.recv()
return data
def distribute(mat):
"""Distributes the mat from root to individual nodes
The data will be distributed along the first axis, as even as possible.
You should make sure that the matrix is in C-contiguous format.
"""
# quick check
if SIZE == 1:
return mat
if is_root():
shape = mat.shape[1:]
segments = get_segments(mat.shape[0])
dtype = mat.dtype
else:
shape = None
segments = None
dtype = None
shape = COMM.bcast(shape)
dtype = COMM.bcast(dtype)
segments = COMM.bcast(segments)
if is_root():
if mat.flags['C_CONTIGUOUS'] != True:
logging.warning('Warning: mat is not contiguous.')
mat = np.ascontiguousarray(mat)
for i in range(1,SIZE):
safe_send_matrix(mat[segments[i]:segments[i+1]], dest=i)
data = mat[:segments[1]].copy()
else:
data = np.empty((segments[RANK+1] - segments[RANK],) + shape,
dtype = dtype)
safe_recv_matrix(data)
return data
def distribute(mat):
"""Distributes the mat from root to individual nodes
The data will be distributed along the first axis, as even as possible.
You should make sure that the matrix is in C-contiguous format.
"""
# quick check
if SIZE == 1:
return mat
if is_root():
shape = mat.shape[1:]
segments = get_segments(mat.shape[0])
dtype = mat.dtype
else:
shape = None
segments = None
dtype = None
shape = COMM.bcast(shape)
dtype = COMM.bcast(dtype)
segments = COMM.bcast(segments)
if is_root():
if mat.flags['C_CONTIGUOUS'] != True:
logging.warning('Warning: mat is not contiguous.')
mat = np.ascontiguousarray(mat)
for i in range(1, SIZE):
safe_send_matrix(mat[segments[i]:segments[i + 1]], dest=i)
data = mat[:segments[1]].copy()
else:
data = np.empty((segments[RANK + 1] - segments[RANK], ) + shape,
dtype=dtype)
safe_recv_matrix(data)
return data
def load_matrix_multi(filename, N = None):
"""Loads the matrix previously dumped by dump_matrix_multi. The MPI size
might be different. The stored files are in the format
filename-xxxxx-of-xxxxx, which we obtain using glob.
Input:
N: (optional) if given, specify the number of parts the matrix is
separated too. Otherwise, the number is automatically inferred by
listing all the files using regexp matching.
"""
files= glob.glob('%s-?????-of-?????.npy' % (filename))
N = len(files)
logging.debug("Loading the matrix from %d parts" % N)
# we will load the length of the data, and then try to distribute them
# as even as possible.
if RANK == 0:
# the root will first taste each file
sizes = np.array([np.load('%s-%05d-of-%05d.npy' % (filename, i, N),
mmap_mode='r').shape[0]
for i in range(N)])
temp = np.load('%s-%05d-of-%05d.npy' % (filename, 0, N),
mmap_mode='r')
shape = temp.shape[1:]
dtype = temp.dtype
else:
sizes = None
shape = None
dtype = None
barrier()
sizes = COMM.bcast(sizes)
shape = COMM.bcast(shape)
dtype = COMM.bcast(dtype)
total = sizes.sum()
segments = get_segments(total)
# now, each node opens the file that overlaps with its data, and reads
# the contents.
my_start = segments[RANK]
my_end = segments[RANK+1]
my_size = my_end - my_start
mat = np.empty((my_size,) + shape, dtype = dtype)
mat = np.empty((my_size,) + shape)
f_start = 0
f_end = 0
for i, size in enumerate(sizes):
f_end += size
if f_start < my_end and f_end > my_start:
file_mat = np.load('%s-%05d-of-%05d.npy' % (filename, i, N),
mmap_mode='r')
mat[max(f_start - my_start, 0):\
min(f_end - my_start, my_size)] = \
file_mat[max(my_start - f_start,0):\
min(my_end - f_start, size)]
f_start += size
return mat
def load_matrix_multi(filename, N=None):
"""Loads the matrix previously dumped by dump_matrix_multi. The MPI size
might be different. The stored files are in the format
filename-xxxxx-of-xxxxx, which we obtain using glob.
Input:
N: (optional) if given, specify the number of parts the matrix is
separated too. Otherwise, the number is automatically inferred by
listing all the files using regexp matching.
"""
files = glob.glob('%s-?????-of-?????.npy' % (filename))
N = len(files)
logging.debug("Loading the matrix from %d parts" % N)
# we will load the length of the data, and then try to distribute them
# as even as possible.
if RANK == 0:
# the root will first taste each file
sizes = np.array([
np.load('%s-%05d-of-%05d.npy' % (filename, i, N),
mmap_mode='r').shape[0] for i in range(N)
])
temp = np.load('%s-%05d-of-%05d.npy' % (filename, 0, N), mmap_mode='r')
shape = temp.shape[1:]
dtype = temp.dtype
else:
sizes = None
shape = None
dtype = None
barrier()
sizes = COMM.bcast(sizes)
shape = COMM.bcast(shape)
dtype = COMM.bcast(dtype)
total = sizes.sum()
segments = get_segments(total)
# now, each node opens the file that overlaps with its data, and reads
# the contents.
my_start = segments[RANK]
my_end = segments[RANK + 1]
my_size = my_end - my_start
mat = np.empty((my_size, ) + shape, dtype=dtype)
mat = np.empty((my_size, ) + shape)
f_start = 0
f_end = 0
for i, size in enumerate(sizes):
f_end += size
if f_start < my_end and f_end > my_start:
file_mat = np.load('%s-%05d-of-%05d.npy' % (filename, i, N),
mmap_mode='r')
mat[max(f_start - my_start, 0):\
min(f_end - my_start, my_size)] = \
file_mat[max(my_start - f_start,0):\
min(my_end - f_start, size)]
f_start += size
return mat
def safe_send_matrix(mat, dest=0, tag=0):
"""A safe send that deals with the mpi4py 2GB limit. should be paired with
safe_recv_matrix. The input mat should be C_CONTIGUOUS. To be safe, we send
the matrix in 1GB chunks.
"""
num_batches = int((mat.nbytes - 1) / _MPI_BUFFER_LIMIT + 1)
if num_batches == 1:
COMM.Send(mat, dest, tag)
else:
logging.debug("The buffer is larger than 1GB, sending in chunks...")
batch_size = int(mat.shape[0] / num_batches)
for i in range(num_batches):
COMM.Send(mat[batch_size * i:batch_size * (i + 1)], dest, tag)
# send the remaining part
if mat.shape[0] > batch_size * num_batches:
COMM.Send(mat[batch_size * num_batches:], dest, tag)
def dump_matrix(mat, filename):
"""Dumps the matrix distributed over machines to one single file.
We do NOT recommend using this - it causes a lot of communications since
all data need to be transferred to root before writing to disk. Instead,
use dump_matrix_multi which stores the matrix to multiple chunks.
"""
if SIZE == 1:
with open(filename, 'w') as fid:
np.save(fid, mat)
else:
mat_sizes = COMM.gather(mat.shape[0])
if is_root():
total_size = sum(mat_sizes)
mat_reduced = np.empty((total_size, ) + mat.shape[1:],
dtype=mat.dtype)
start = mat_sizes[0]
mat_reduced[:start] = mat
for i in range(1, SIZE):
safe_recv_matrix(mat_reduced[start:start + mat_sizes[i]],
source=i)
start += mat_sizes[i]
with open(filename, 'w') as fid:
np.save(fid, mat_reduced)
else:
safe_send_matrix(mat, dest=0)
barrier()
def dump_matrix(mat, filename):
"""Dumps the matrix distributed over machines to one single file.
We do NOT recommend using this - it causes a lot of communications since
all data need to be transferred to root before writing to disk. Instead,
use dump_matrix_multi which stores the matrix to multiple chunks.
"""
if SIZE == 1:
with open(filename,'w') as fid:
np.save(fid, mat)
else:
mat_sizes = COMM.gather(mat.shape[0])
if is_root():
total_size = sum(mat_sizes)
mat_reduced = np.empty((total_size,) + mat.shape[1:],
dtype = mat.dtype)
start = mat_sizes[0]
mat_reduced[:start] = mat
for i in range(1,SIZE):
safe_recv_matrix(mat_reduced[start:start+mat_sizes[i]],
source = i)
start += mat_sizes[i]
with open(filename,'w') as fid:
np.save(fid, mat_reduced)
else:
safe_send_matrix(mat, dest = 0)
barrier()
def barrier(tag=0, sleep=0.01):
''' A better mpi barrier
The original MPI.comm.barrier() may cause idle processes to still occupy
the CPU, while this barrier waits.
'''
if SIZE == 1:
return
mask = 1
while mask < SIZE:
dst = (RANK + mask) % SIZE
src = (RANK - mask + SIZE) % SIZE
req = COMM.isend(None, dst, tag)
while not COMM.Iprobe(src, tag):
time.sleep(sleep)
COMM.recv(None, src, tag)
req.Wait()
mask <<= 1
def elect():
'''elect() randomly chooses a node from all the nodes as the president.
Input:
None
Output:
the rank of the president
'''
president = COMM.bcast(np.random.randint(SIZE))
return president
def barrier(tag=0, sleep=0.01):
''' A better mpi barrier
The original MPI.comm.barrier() may cause idle processes to still occupy
the CPU, while this barrier waits.
'''
if SIZE == 1:
return
mask = 1
while mask < SIZE:
dst = (RANK + mask) % SIZE
src = (RANK - mask + SIZE) % SIZE
req = COMM.isend(None, dst, tag)
while not COMM.Iprobe(src, tag):
time.sleep(sleep)
COMM.recv(None, src, tag)
req.Wait()
mask <<= 1
def elect():
'''elect() randomly chooses a node from all the nodes as the president.
Input:
None
Output:
the rank of the president
'''
president = COMM.bcast(np.random.randint(SIZE))
return president
# MPI
try:
from mpi4py import MPI
COMM = MPI.COMM_WORLD
except Exception, e:
sys.stderr.write(\
"Warning: I cannot import mpi4py. Using a dummpy single noded "\
"implementation instead. The program will run in single node mode "\
"even if you executed me with mpirun or mpiexec.\n")
sys.stderr.write("We STRONGLY recommend you to try to install mpi and "\
"mpi4py.\n")
sys.stderr.write("mpi4py exception message is:")
sys.stderr.write(repr(Exception) + repr(e))
from _mpi_dummy import COMM
RANK = COMM.Get_rank()
SIZE = COMM.Get_size()
_HOST_RAW = socket.gethostname()
# this is the hack that removes things like ".icsi.berkeley.edu"
if _HOST_RAW.find('.') == -1:
HOST = _HOST_RAW
else:
HOST = _HOST_RAW[:_HOST_RAW.find('.')]
_MPI_PRINT_MESSAGE_TAG = 560710
_MPI_BUFFER_LIMIT = 1073741824
# we need to set the random seed different for each mpi instance
random.seed(time.time() * RANK)
def mkdir(dirname):
def agree(decision):
"""agree() makes the decision consistent by propagating the decision of the
root to everyone
"""
return COMM.bcast(decision)
def load_matrix_multi(filename, N=None, name=None):
"""Loads the matrix previously dumped by dump_matrix_multi. The MPI size
might be different. The stored files are in the format
filename-xxxxx-of-xxxxx, which we obtain using glob.
Input:
name: if the input is a hdf5 mat file, specify the name here.
"""
if type(filename) is str:
# we use our default format
files = glob.glob('%s-?????-of-?????.npy' % (filename))
files.sort()
else:
files = list(filename)
N = len(files)
logging.debug("Loading the matrix from %d parts" % N)
# we will load the length of the data, and then try to distribute them
# as even as possible.
if RANK == 0:
# the root will first taste each file
if files[0][-3:] == 'npy':
sizes = np.array(
[np.load(f, mmap_mode='r').shape[0] for f in files])
temp = np.load(files[0], mmap_mode='r')
shape = temp.shape[1:]
dtype = temp.dtype
elif files[0][-3:] == 'mat':
sizes = []
for f in files:
fid = h5py.File(f, 'r')
sizes.append(fid[name].shape[0])
shape = fid[name].shape[1:]
dtype = fid[name].dtype
fid.close()
sizes = np.array(sizes)
else:
sizes = None
shape = None
dtype = None
barrier()
sizes = COMM.bcast(sizes)
shape = COMM.bcast(shape)
dtype = COMM.bcast(dtype)
total = sizes.sum()
segments = get_segments(total)
# now, each node opens the file that overlaps with its data, and reads
# the contents.
my_start = segments[RANK]
my_end = segments[RANK + 1]
my_size = my_end - my_start
mat = np.empty((my_size, ) + shape, dtype=dtype)
mat = np.empty((my_size, ) + shape)
f_start = 0
f_end = 0
for i, size in enumerate(sizes):
f_end += size
if f_start < my_end and f_end > my_start:
if files[i][-3:] == 'npy':
file_mat = np.load(files[i], mmap_mode='r')
mat[max(f_start - my_start, 0):\
min(f_end - my_start, my_size)] = \
file_mat[max(my_start - f_start,0):\
min(my_end - f_start, size)]
elif files[i][-3:] == 'mat':
fid = h5py.File(files[i], 'r')
mat[max(f_start - my_start, 0):\
min(f_end - my_start, my_size)] = \
fid[name][max(my_start - f_start,0):\
min(my_end - f_start, size)]
fid.close()
f_start += size
return mat
def load_matrix_multi(filename, N = None, name=None):
"""Loads the matrix previously dumped by dump_matrix_multi. The MPI size
might be different. The stored files are in the format
filename-xxxxx-of-xxxxx, which we obtain using glob.
Input:
name: if the input is a hdf5 mat file, specify the name here.
"""
if type(filename) is str:
# we use our default format
files = glob.glob('%s-?????-of-?????.npy' % (filename))
files.sort()
else:
files = list(filename)
N = len(files)
logging.debug("Loading the matrix from %d parts" % N)
# we will load the length of the data, and then try to distribute them
# as even as possible.
if RANK == 0:
# the root will first taste each file
if files[0][-3:] == 'npy':
sizes = np.array([np.load(f, mmap_mode='r').shape[0]
for f in files])
temp = np.load(files[0], mmap_mode='r')
shape = temp.shape[1:]
dtype = temp.dtype
elif files[0][-3:] == 'mat':
sizes = []
for f in files:
fid = h5py.File(f, 'r')
sizes.append(fid[name].shape[0])
shape = fid[name].shape[1:]
dtype = fid[name].dtype
fid.close()
sizes = np.array(sizes)
else:
sizes = None
shape = None
dtype = None
barrier()
sizes = COMM.bcast(sizes)
shape = COMM.bcast(shape)
dtype = COMM.bcast(dtype)
total = sizes.sum()
segments = get_segments(total)
# now, each node opens the file that overlaps with its data, and reads
# the contents.
my_start = segments[RANK]
my_end = segments[RANK+1]
my_size = my_end - my_start
mat = np.empty((my_size,) + shape, dtype = dtype)
mat = np.empty((my_size,) + shape)
f_start = 0
f_end = 0
for i, size in enumerate(sizes):
f_end += size
if f_start < my_end and f_end > my_start:
if files[i][-3:] == 'npy':
file_mat = np.load(files[i], mmap_mode='r')
mat[max(f_start - my_start, 0):\
min(f_end - my_start, my_size)] = \
file_mat[max(my_start - f_start,0):\
min(my_end - f_start, size)]
elif files[i][-3:] == 'mat':
fid = h5py.File(files[i], 'r')
mat[max(f_start - my_start, 0):\
min(f_end - my_start, my_size)] = \
fid[name][max(my_start - f_start,0):\
min(my_end - f_start, size)]
fid.close()
f_start += size
return mat
def agree(decision):
"""agree() makes the decision consistent by propagating the decision of the
root to everyone
"""
return COMM.bcast(decision)
