def persist(self) -> Union[xa.DataArray,DataArrayGroupBy]: from dask.distributed import Client xrd: xa.DataArray = self.xr if isinstance(xrd,DataArrayGroupBy): self.logger.warn( " EDASArray.persist returning DataArrayGroupBy" ) return xrd if self.loaded_data is None: client = Client.current() if client is None: self.loaded_data = xrd.load().persist() else: client.persist( xrd ) self.loaded_data = xrd return self.loaded_data
def get_client( address=None, auto_spawn=True, worker_log_level="ERROR", **clustser_kwargs, ): """ Args: address (str, optional): address of the cluster scheduler, or 'slurm' to launch a dask cluster through SLURM auto_spawn (bool, optional): automagically spawn cluster if not found work_log_level (str, optional): worker log level """ cluster_klass, client = None, None if address == "slurm": # create SLURM jobs cluster_klass = ManagedSLURMCluster elif address is None: # nothing specified, use: # - already connected client # - spawn new local cluster try: # we try to acquire current session first client = Client.current() address = client.scheduler_info()["address"] logger.info(f"connect to existing cluster (scheduler: {address})") yield client # NOTE we do NOT close client when using this method, managed by others except ValueError: # nothing exists, continue to spawn managed cluster if not auto_spawn: raise RuntimeError("please spawn a dask cluster first") cluster_klass = ManagedLocalCluster # local cluster needs address info clustser_kwargs.update({"address": address}) else: # directly specify the scheduler to connect to client = Client(address) yield client client.close() # NOTE we open this client, therefore, we need to close it ourself if not cluster_klass: # no need to spawn a cluster return with cluster_klass(**clustser_kwargs) as cluster: client = cluster.client # register loggers try: import coloredlogs except ImportError: logger.install( "install `coloredlogs` to configure loggers automatically") else: def install_logger(dask_worker): # we know this is annoying, silence it logging.getLogger("tifffile").setLevel(logging.ERROR) coloredlogs.install( level=worker_log_level, fmt="%(asctime)s %(levelname)s %(message)s", datefmt="%H:%M:%S", ) logger.debug( f'install logger for workers, level="{worker_log_level}"') client.register_worker_callbacks(install_logger) yield client
if slurm: cluster = dask_jobqueue.SLURMCluster( queue=os.environ['SLURM_JOB_PARTITION'], cores=1, walltime='0', memory=opts['memory'], local_directory=os.getenv('TMPDIR', '/tmp')) client = Client(cluster) cluster.start_workers(opts['ntasks']) else: cluster = LocalCluster(n_workers=opts['ntasks'], threads_per_worker=1) client = Client(cluster) print(client.current().cluster.scheduler) # read model configuration parameters = configurate() # Sterope Main Algorithm population = populate() # simulate levels simulate() # evaluate sensitivity sensitivity = evaluate() # write reports report() # move and organize results backup()
def scheduler_info_solid(context): with context.resources.dask.client.as_current(): client = Client.current() yield Output(client.scheduler_info(), "scheduler_info") yield Output(client.nthreads(), "nthreads")