Python dask.distributed.Client() Examples

def LaunchDaskDistributedClient(self, scheduler_ip=None, scheduler_port=None):

        if self.parallel and self.parallel_model == "dask" and self.is_dask_scheduler_initialised is False:

            from multiprocessing.pool import ThreadPool
            try:
                import dask
                from dask.distributed import Client, LocalCluster
            except ImportError:
                raise ImportError("dask is not installed. Install it 'using pip install dask[complete]'")

            dask.config.set(pool=ThreadPool(self.no_of_cpu_cores))
            # INITIALISE CLUSTER
            if scheduler_ip is None:
                cluster = LocalCluster(n_workers=self.no_of_cpu_cores, processes=False, threads_per_worker=None)
                client = Client(cluster)
            else:
                client = Client(scheduler_ip)

            self.dask_client = client

            self.is_dask_scheduler_initialised = True 

def setup_dask(scheduler, retries=-1):
    if scheduler is None or scheduler == "{scheduler}":
        print("Setting up local cluster...")
        return Client()
    succeeded = False
    try_num = 0
    while not succeeded:
        try_num += 1
        if try_num == retries:
            raise Exception("Failed to connect to Dask client")
        try:
            client = Client(scheduler, timeout=60)
            succeeded = True
        except Exception as e:  # pylint: disable=broad-except
            print(e)
        time.sleep(15)

    return client 

def setup_dask(scheduler, retries=-1):
    if scheduler is None or scheduler == "{scheduler}":
        print("Setting up local cluster...")
        return Client()
    succeeded = False
    try_num = 0
    while not succeeded:
        try_num += 1
        if try_num == retries:
            raise Exception("Failed to connect to Dask client")
        try:
            client = Client(scheduler, timeout=60)
            succeeded = True
        except Exception as e:  # pylint: disable=broad-except
            print(e)
        time.sleep(15)

    return client 

def dask_client_create(**kwargs):
    """
    Create Dask client object. The function is trivial and introduced so that
    Dask client is created in uniform way throughout the program.

    Parameters
    ----------
    kwargs: dict, optional
        kwargs will be passed to the Dask client constructor

    Returns
    -------
    client: dask.distributed.Client
        Dask client object
    """
    _kwargs = {"processes": True, "silence_logs": logging.ERROR}
    _kwargs.update(kwargs)
    client = Client(**_kwargs)
    dask.config.set(shuffle="disk")
    path_dask_data = os.path.expanduser("~/.dask")
    dask.config.set({"temporary_directory": path_dask_data})
    return client 

def run_search():

    from dask.distributed import Client, LocalCluster
    import joblib
    import hypertunity as ht

    #client = Client(scheduler_file='scheduler.json')
    client = Client()
    print(client)

    domain = ht.Domain({
                    "cost_rate": set([-.8])
    })

    # with joblib.parallel_backend('dask'):
    #     with joblib.Parallel() as parallel:
    #         print("Doing the work ... ")
    #         results = parallel(joblib.delayed(run_games)(*domain.sample().as_namedtuple()) for s in range(1))
    #
    # print(results)
    run_games(-.8) 

def activate_client(self,
                        LSF = True,
                        num_processes = 2,
                        adapt = False):

        if LSF:
            from dask_jobqueue import LSFCluster
            cluster = LSFCluster()
            self._adapt = adapt
            self.num_processes = num_processes

            if self._adapt:
                _logger.debug(f"adapting cluster from 1 to {self.num_processes} processes")
                cluster.adapt(minimum = 2, maximum = self.num_processes, interval = "1s")
            else:
                _logger.debug(f"scaling cluster to {self.num_processes} processes")
                cluster.scale(self.num_processes)

            _logger.debug(f"scheduling cluster with client")
            self.client = distributed.Client(cluster)
        else:
            self.client = None
            self._adapt = False
            self.num_processes = 0 

def main():
    client = Client()  # noqa

    categories = ["category_%d" % i for i in range(26)]
    columns = ["click"] + ["numeric_%d" % i for i in range(13)] + categories

    df = dd.read_csv("day_1", sep="\t", names=columns, header=None)

    encoding = {c: "bytes" for c in categories}
    fixed = {c: 8 for c in categories}
    df.to_parquet(
        "day-1-bytes.parquet",
        object_encoding=encoding,
        fixed_text=fixed,
        compression="SNAPPY",
    ) 

def client():
    from dask.distributed import Client
    with Client(n_workers=4) as dask_client:
        yield dask_client 

def test_env(pod_spec, ns, auth):
    async with KubeCluster(
        pod_spec, env={"ABC": "DEF"}, namespace=ns, auth=auth, **cluster_kwargs
    ) as cluster:
        cluster.scale(1)
        await cluster
        async with Client(cluster, asynchronous=True) as client:
            while not cluster.scheduler_info["workers"]:
                await asyncio.sleep(0.1)
            env = await client.run(lambda: dict(os.environ))
            assert all(v["ABC"] == "DEF" for v in env.values()) 

def test_pod_from_minimal_dict(image_name, ns, auth):
    spec = {
        "spec": {
            "containers": [
                {
                    "args": [
                        "dask-worker",
                        "$(DASK_SCHEDULER_ADDRESS)",
                        "--nthreads",
                        "1",
                        "--death-timeout",
                        "60",
                    ],
                    "command": None,
                    "image": image_name,
                    "imagePullPolicy": "IfNotPresent",
                    "name": "worker",
                }
            ]
        }
    }

    async with KubeCluster.from_dict(
        spec, namespace=ns, auth=auth, **cluster_kwargs
    ) as cluster:
        cluster.adapt()
        async with Client(cluster, asynchronous=True) as client:
            future = client.submit(lambda x: x + 1, 10)
            result = await future
            assert result == 11 

def test_start_with_workers(pod_spec, ns, auth):
    async with KubeCluster(
        pod_spec, n_workers=2, namespace=ns, auth=auth, **cluster_kwargs
    ) as cluster:
        async with Client(cluster, asynchronous=True) as client:
            while len(cluster.scheduler_info["workers"]) != 2:
                await asyncio.sleep(0.1) 

def client(cluster):
    with Client(cluster) as client:
        yield client 

def test_env(pod_spec, loop, ns):
    with KubeCluster(pod_spec, env={"ABC": "DEF"}, loop=loop, namespace=ns) as cluster:
        cluster.scale(1)
        with Client(cluster, loop=loop) as client:
            while not cluster.scheduler_info["workers"]:
                sleep(0.1)
            env = client.run(lambda: dict(os.environ))
            assert all(v["ABC"] == "DEF" for v in env.values()) 

def test_pod_from_minimal_dict(image_name, loop, ns):
    spec = {
        "spec": {
            "containers": [
                {
                    "args": [
                        "dask-worker",
                        "$(DASK_SCHEDULER_ADDRESS)",
                        "--nthreads",
                        "1",
                        "--death-timeout",
                        "60",
                    ],
                    "command": None,
                    "image": image_name,
                    "imagePullPolicy": "IfNotPresent",
                    "name": "worker",
                }
            ]
        }
    }

    with KubeCluster.from_dict(spec, loop=loop, namespace=ns) as cluster:
        cluster.adapt()
        with Client(cluster, loop=loop) as client:
            future = client.submit(lambda x: x + 1, 10)
            result = future.result()
            assert result == 11 

def test_to_keys_numpy_array():
    rng = np.random.RandomState(0)
    arr = rng.randn(20, 30)
    df = pd.DataFrame(data=arr)
    dsk = {}
    grid_search_keys = list(dms.utils.to_keys(dsk, arr, df))
    with Client() as client:
        data_futures = client.scatter([arr, df])
    assert grid_search_keys == [f.key for f in data_futures] 

def client():
    from dask.distributed import Client
    with Client(n_workers=4) as dask_client:
        yield dask_client 

def client():
    from dask.distributed import Client
    with Client(n_workers=4) as dask_client:
        yield dask_client 

def client():
    from dask.distributed import Client
    with Client(n_workers=4) as dask_client:
        yield dask_client 

def client(cluster):
    async with Client(cluster, asynchronous=True) as client:
        yield client 

def __init__(self, address: str,
                 tls_ca_file: str = None, tls_client_cert: str = None, tls_client_key: str = None,
                 require_encryption: bool = False):
        # authentication
        sec = None
        if tls_ca_file and tls_client_cert and tls_client_key:
            sec = Security(tls_ca_file=tls_ca_file,
                           tls_client_cert=tls_client_cert,
                           tls_client_key=tls_client_key,
                           require_encryption=require_encryption)

        # init
        self._client = Client(address=address, security=sec)
        self._client.register_worker_callbacks(Remote._worker_startup) 

def scheduler(request):
    if request.param.startswith("distributed"):
        kw = {"dashboard_address": None}

        if request.param == "distributed-threads":
            kw["processes"] = False

        client = Client(**kw)
        yield client
        client.close()

    else:
        yield request.param 

def client():
    client = Client(processes=False, dashboard_address=None)
    yield client
    client.close() 

def __init__(self, cluster_type):
        if cluster_type not in ALLOWED_CLUSTER_TYPES:
            raise Exception("Can't choose this type of cluster for now. Choose one from: {}".format(
                ALLOWED_CLUSTER_TYPES))
        self.cluster_type = cluster_type
        self._cluster = self.get_cluster()
        self._client = Client(self._cluster) 

def _get_kubernetes_cluster(worker_template_path=WORKER_TEMPLATE_PATH):
        from dask_kubernetes import KubeCluster

        cluster = KubeCluster.from_yaml(worker_template_path)
        return Client(cluster) 

def client(self):
        if self._client is None:
            self._client = Client(self._cluster)
        return self._client 

def client(self):
        """
        Get the Client associated with this job. The Client connects users to a dask.distributed compute cluster. It provides an asynchronous user interface around functions and futures.

        Returns:
            dask.distributed.Client: Client initialized with Scheduler configuration if defined, otherwise bound locally.
        """  # noqa: #501
        if not getattr(self, '_client', None):
            if self.heartbeat_interval == 0:
                heartbeat_interval = None
            else:
                heartbeat_interval = self.heartbeat_interval

            if self.timeout == 0:
                timeout = '__no_default__'
            else:
                timeout = self.timeout

            try:
                # validating the invalid scheduler
                client = Client(
                    address=self.host,
                    heartbeat_interval=heartbeat_interval,
                    timeout=timeout
                )

            except Exception:
                log.exception('Dask Client Init Error')
                raise DaskJobException('Invalid scheduler is provided')

            self._client = client

        return self._client 

def client(self):

        if self.scheduler:
            return self.scheduler.client
        else:
            return Client() 

def _execute(self, future_or_delayed, *args, **kwargs):
        """
        Execute Delayed jobs using the distributed Client to get a future object. Save the key of the future object for later use.

        Args:
            future_or_delayed (dask.Delayed or dask.distributed.Future): dask task object to track using TethysJobs.
        """  # noqa: E501
        if not isinstance(future_or_delayed, Future) and not isinstance(future_or_delayed, Delayed):
            raise ValueError('Must pass a valid instance of Delayed or Future.')

        if isinstance(future_or_delayed, Delayed):
            future = self.client.compute(future_or_delayed)
        else:
            future = future_or_delayed

        self.key = future.key

        # NOTE: Job may not actually be running at this point, but we don't have another
        # way to know when the first node in the workflow starts running.
        self._status = 'RUN'
        self.start_time = timezone.now()

        # Send key to the dask scheduler so the scheduler knows which jobs to send status updates to Tethys.
        self.client.set_metadata(self.key, True)

        # Save updated attributes
        self.save()

        # Must use fire and forget to ensure job runs after the future goes out of scope.
        fire_and_forget(future)

        # Save this client to close it after obtaining the result.
        global client_fire_forget
        client_fire_forget = self.client 

def test_cluster_create(cluster):
    cluster.scale(1)
    await cluster
    async with Client(cluster, asynchronous=True) as client:
        result = await client.submit(lambda x: x + 1, 10)
        assert result == 11 

def test_dask_multiprocessing(tmpdir):
    """
    Test that dask multiprocessing works on Python 3.
    """
    # Command to start the kernel
    cmd = "from spyder_kernels.console import start; start.main()"

    with setup_kernel(cmd) as client:
        # Remove all variables
        client.execute("%reset -f")
        client.get_shell_msg(block=True, timeout=TIMEOUT)

        # Write multiprocessing code to a file
        # Runs two times to verify that in the second case it doesn't break
        code = """
from dask.distributed import Client

if __name__=='__main__':
    client = Client()
    client.close()
    x = 'hello'
"""
        p = tmpdir.join("mp-test.py")
        p.write(code)

        # Run code two times
        client.execute("runfile(r'{}')".format(to_text_string(p)))
        client.get_shell_msg(block=True, timeout=TIMEOUT)

        client.execute("runfile(r'{}')".format(to_text_string(p)))
        client.get_shell_msg(block=True, timeout=TIMEOUT)

        # Verify that the `x` variable is defined
        client.inspect('x')
        msg = client.get_shell_msg(block=True, timeout=TIMEOUT)
        content = msg['content']
        assert content['found']