Python distributed.LocalCluster() Examples

def spawn_cluster_and_client(
    address: Optional[str] = None, **kwargs
) -> Tuple[Optional[LocalCluster], Optional[Client]]:
    """
    If provided an address, create a Dask Client connection.
    If not provided an address, create a LocalCluster and Client connection.
    If not provided an address, other Dask kwargs are accepted and passed down to the
    LocalCluster object.

    Notes
    -----
    When using this function, the processing machine or container must have networking
    capabilities enabled to function properly.
    """
    cluster = None
    if address is not None:
        client = Client(address)
        log.info(f"Connected to Remote Dask Cluster: {client}")
    else:
        cluster = LocalCluster(**kwargs)
        client = Client(cluster)
        log.info(f"Connected to Local Dask Cluster: {client}")

    return cluster, client 

def shutdown_cluster_and_client(
    cluster: Optional[LocalCluster], client: Optional[Client]
) -> Tuple[Optional[LocalCluster], Optional[Client]]:
    """
    Shutdown a cluster and client.

    Notes
    -----
    When using this function, the processing machine or container must have networking
    capabilities enabled to function properly.
    """
    if cluster is not None:
        cluster.close()
    if client is not None:
        client.shutdown()
        client.close()

    return cluster, client 

def cluster_and_client(address: Optional[str] = None, **kwargs):
    """
    If provided an address, create a Dask Client connection.
    If not provided an address, create a LocalCluster and Client connection.
    If not provided an address, other Dask kwargs are accepted and passed down to the
    LocalCluster object.

    These objects will only live for the duration of this context manager.

    Examples
    --------
    >>> with cluster_and_client() as (cluster, client):
    ...     img1 = AICSImage("1.tiff")
    ...     img2 = AICSImage("2.czi")
    ...     other processing

    Notes
    -----
    When using this context manager, the processing machine or container must have
    networking capabilities enabled to function properly.
    """
    try:
        cluster, client = spawn_cluster_and_client(address=address, **kwargs)
        yield cluster, client
    finally:
        shutdown_cluster_and_client(cluster=cluster, client=client) 

def start_cluster(diagnostics_port=0):
    "Set up a LocalCluster for distributed"
    
    hostname = socket.gethostname()
    n_workers = os.cpu_count() // 2
    cluster = LocalCluster(ip='localhost',
                       n_workers=n_workers,
                       diagnostics_port=diagnostics_port,
                       memory_limit=6e9)
    client = Client(cluster)

    params = { 'bokeh_port': cluster.scheduler.services['bokeh'].port,
           'user': getpass.getuser(),
           'scheduler_ip': cluster.scheduler.ip,
           'hostname': hostname, }

    print("If the link to the dashboard below doesn't work, run this command on a local terminal to set up a SSH tunnel:")
    print()
    print("  ssh -N -L {bokeh_port}:{scheduler_ip}:{bokeh_port} {hostname}.nci.org.au -l {user}".format(**params) )
    
    return client 

def _n_workers_for_local_cluster(calcs):
    """The number of workers used in a LocalCluster

    An upper bound is set at the cpu_count or the number of calcs submitted,
    depending on which is smaller.  This is to prevent more workers from
    being started than needed (but also to prevent too many workers from
    being started in the case that a large number of calcs are submitted).
    """
    return min(cpu_count(), len(calcs)) 

def _exec_calcs(calcs, parallelize=False, client=None, **compute_kwargs):
    """Execute the given calculations.

    Parameters
    ----------
    calcs : Sequence of ``aospy.Calc`` objects
    parallelize : bool, default False
        Whether to submit the calculations in parallel or not
    client : distributed.Client or None
        The distributed Client used if parallelize is set to True; if None
        a distributed LocalCluster is used.
    compute_kwargs : dict of keyword arguments passed to ``Calc.compute``

    Returns
    -------
    A list of the values returned by each Calc object that was executed.
    """
    if parallelize:
        def func(calc):
            """Wrap _compute_or_skip_on_error to require only the calc
            argument"""
            if 'write_to_tar' in compute_kwargs:
                compute_kwargs['write_to_tar'] = False
            return _compute_or_skip_on_error(calc, compute_kwargs)

        if client is None:
            n_workers = _n_workers_for_local_cluster(calcs)
            with distributed.LocalCluster(n_workers=n_workers) as cluster:
                with distributed.Client(cluster) as client:
                    result = _submit_calcs_on_client(calcs, client, func)
        else:
            result = _submit_calcs_on_client(calcs, client, func)
        if compute_kwargs['write_to_tar']:
            _serial_write_to_tar(calcs)
        return result
    else:
        return [_compute_or_skip_on_error(calc, compute_kwargs)
                for calc in calcs] 

def external_client():
    # Explicitly specify we want only 4 workers so that when running on
    # continuous integration we don't request too many.
    cluster = distributed.LocalCluster(n_workers=4)
    client = distributed.Client(cluster)
    yield client
    client.close()
    cluster.close() 

def setUp(self):
        self.dagbag = DagBag(include_examples=True)
        self.cluster = LocalCluster() 

def setup():
    from distributed import LocalCluster, Client
    cluster = LocalCluster(n_workers=1, threads_per_worker=1, processes=False)
    use_distributed(Client(cluster)) 

def _get_local_cluster():
        # TODO: Add more parameters and configurations.
        from distributed import LocalCluster
        return LocalCluster()