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()