def get_cluster(cluster_type, **kwargs):
"""Generic dask cluster wrapper"""
# check input cluster type
cluster_type = cluster_type.lower()
cluster_list = ['lsf','pbs','slurm','local']
if cluster_type not in cluster_list:
msg = "Cluster type '{}' not supported".format(cluster_type)
msg += '\nsupported cluster types: {}'.format(cluster_list)
raise ValueError(msg)
print("Dask cluster type: {}".format(cluster_type))
# No need to do the extra configuration checking if using LocalCluster
if cluster_type == 'local':
return LocalCluster()
# check input config name
if 'config_name' in kwargs.keys():
kwargs['config_name'] = check_config_name(kwargs['config_name'], cluster_type)
print("Dask config name: {}".format(kwargs['config_name']))
# check walltime format for each cluster type
if 'walltime' in kwargs.keys():
kwargs['walltime'] = check_walltime_format(kwargs["walltime"], cluster_type)
print('Dask worker walltime: {}'.format(kwargs['walltime']))
# initiate cluster object
if cluster_type == 'lsf':
cluster = LSFCluster(**kwargs)
elif cluster_type == 'pbs':
cluster = PBSCluster(**kwargs)
elif cluster_type == 'slurm':
cluster = SLURMCluster(**kwargs)
# Print and write job command file for HPC cluster types
print("JOB COMMAND CALLED FROM PYTHON:\n\n", cluster.job_script())
with open('dask_command_run_from_python.txt', 'w') as f:
f.write(cluster.job_script() + '\n')
return cluster
return total
if __name__ == "__main__":
cluster = LSFCluster(
name='worker_bee',
queue='general', # the queue on Pegasus
project='insarlab', # your project name
cores=2,
memory='2GB', # unused by Pegasus but a required param
walltime='00:30', # how long the worker will run for
interface='ib0', # which network to use. NECESSARY PARAM
job_extra=[
'-R "rusage[mem=2500]"', # how to actually define memory usage
"-o WORKER-%J.out"
], # where to write worker output files
python=sys.executable, # Where to look for Python executable
config_name='lsf') # define your own config in a .yaml file
cluster.scale(20)
print("JOB FILE:", cluster.job_script())
client = Client(cluster)
print("Time to run sequential code:", timeit(stmt=sequential_main,
number=1))
print("Time to run parallel code:", timeit(stmt=distributed_main,
number=1))
print("Time to run parallel code with ~0 data transfer:",
timeit(stmt=distributed_main2, number=1))
client.close()