Source code for lenstools.utils.decorators

import sys
from .mpi import MPIWhirlPool

try:
	from mpi4py import MPI
except ImportError:
	MPI = None

class Parallelize(object):

	@classmethod
	def masterworker(cls,func):

		def spreaded_func(*args,**kwargs):

			#MPI Pool
			try:
				pool = MPIWhirlPool()
			except ValueError:
				pool = None

			if (pool is not None) and (not pool.is_master()):
				pool.wait()
				pool.comm.Barrier()
				MPI.Finalize()
				sys.exit(0)

			#Replace the pool in the arguments with the newly created one
			if "pool" in kwargs.keys():
				kwargs["pool"] = pool

			#Execute
			result = func(*args,**kwargs)

			#Finish
			if pool is not None:
				pool.close()
				pool.comm.Barrier()
				MPI.Finalize()

			#Return the result
			if pool is None or pool.is_master():
				return result
		
		#Restore the documentation
		spreaded_func.__doc__ = func.__doc__

		#Return
		return spreaded_func