import os
import contextlib
from copy import deepcopy
import functools
import logging
import signal
from typing import Any, Optional, Union, Callable
from collections.abc import Iterable
import uuid

from dask import distributed as dd
import dask

from libertem.common.threading import set_num_threads_env

from .base import BaseJobExecutor, AsyncAdapter, ResourceError
from libertem.common.executor import (
    JobCancelledError, TaskCommHandler, TaskProtocol, Environment, WorkerContext,
)
from libertem.common.async_utils import sync_to_async
from libertem.common.scheduler import Worker, WorkerSet
from libertem.common.backend import set_use_cpu, set_use_cuda
from libertem.common.async_utils import adjust_event_loop_policy
from .utils import assign_cudas

log = logging.getLogger(__name__)


class DaskWorkerContext(WorkerContext):
    def __init__(self, comms_topic: Optional[str]):
        # DaskWorkerContext sends all messages via a single unique topic id
        # which are later unpacked on the client node; this allows us to handle
        # concurrent runs using the same executor with separate comms channels
        self._comms_topic = comms_topic

    @property
    def dask_worker(self):
        try:
            return self._worker
        except AttributeError:
            self._worker = dd.get_worker()
            return self._worker

    def signal(self, ident: str, topic: str, msg_dict: dict[str, Any]):
        if self._comms_topic is None:
            # Scheduler Dask does not have comms so don't send
            return
        msg_dict.update({'ident': ident, 'topic': topic})
        try:
            self.dask_worker.log_event(self._comms_topic, msg_dict)
        except AttributeError:
            # No structured logs available in this Dask
            # Catch the exception here just in case there is
            # version / API mismatch
            pass


@contextlib.contextmanager
def set_worker_log_level(level: Union[str, int], force: bool = False):
    """
    Set the dask.distributed log level for any processes spawned
    within the context manager. If force is False, don't overwrite
    any existing environment variable.
    """
    env_keys = ['DASK_DISTRIBUTED__LOGGING__DISTRIBUTED']
    try:
        old_env = {
            k: os.environ[k]
            for k in env_keys
            if k in os.environ
        }
        os.environ.update({
            k: str(level)
            for k in env_keys
            if force or (k not in old_env)
        })
        yield
    finally:
        os.environ.update(old_env)
        for key in env_keys:
            if key not in old_env:
                del os.environ[key]


def worker_setup(resource, device):
    # Disable handling Ctrl-C on the workers for a local cluster
    # since the nanny restarts workers in that case and that gets mixed
    # with Ctrl-C handling of the main process, at least on Windows
    signal.signal(signal.SIGINT, signal.SIG_IGN)
    if resource == "CUDA":
        set_use_cuda(device)
    elif resource == "CPU":
        set_use_cpu(device)
    else:
        raise ValueError("Unknown resource %s, use 'CUDA' or 'CPU'", resource)


def cluster_spec(
    cpus: Union[int, Iterable[int]],
    cudas: Union[int, Iterable[int]],
    has_cupy: bool,
    name: str = 'default',
    num_service: int = 1,
    options: Optional[dict] = None,
    preload: Optional[tuple[str, ...]] = None
):
    '''
    Create a worker specification dictionary for a LiberTEM Dask cluster

    The return from this function can be passed to :code:`DaskJobExecutor.make_local(spec=spec)`.

    This creates a Dask cluster spec with special initializations and resource tags
    for CPU + GPU processing in LiberTEM.
    See :ref:`cluster spec` for an example.
    See https://distributed.dask.org/en/stable/api.html#distributed.SpecCluster
    for more info on cluster specs.

    Parameters
    ----------
    cpus: int | Iterable[int]
        IDs for CPU workers as an iterable, or an integer number of workers to create.
        Currently no pinning is used, i.e. this specifies the total
        number and identification of workers, not the CPU cores that are used.
    cudas: int | Iterable[int]
        IDs for CUDA device workers as an iterable, or an integer number of GPU workers to
        create. LiberTEM will use the IDs specified or assign round-robin to the available devices.
        In the iterable case these have to match CUDA device IDs on the system.
        Specify the same ID multiple times to spawn multiple workers on the same CUDA device.
    has_cupy: bool
        Specify if the cluster should signal that it supports GPU-based array programming using
        CuPy
    name
        Prefix for the worker names
    num_service
        Number of additional workers that are reserved for service tasks. Computation tasks
        will not be scheduled on these workers, which guarantees responsive behavior for file
        browsing etc.
    options
        Options to pass through to every worker. See Dask documentation for details
    preload
        Items to preload on workers in addition to LiberTEM-internal preloads.
        This can be used to load libraries, for example HDF5 filter plugins before h5py is used.
        See https://docs.dask.org/en/stable/how-to/customize-initialization.html#preload-scripts
        for more information.

    See also
    --------
    :func:`libertem.utils.devices.detect`
    '''
    if options is None:
        options = {}
    if preload is None:
        preload = ()
    if options.get("nthreads") is None:
        options["nthreads"] = 1
    if options.get("silence_logs") is None:
        options["silence_logs"] = logging.WARN

    workers_spec = {}

    cpu_options = deepcopy(options)
    cpu_options["resources"] = {"CPU": 1, "compute": 1, "ndarray": 1}
    cpu_base_spec = {
        "cls": dd.Nanny,
        "options": cpu_options,
    }

    # Service workers not for computation

    service_options = deepcopy(options)
    service_options["resources"] = {}
    service_base_spec = {
        "cls": dd.Nanny,
        "options": service_options
    }

    cuda_options = deepcopy(options)
    cuda_options["resources"] = {"CUDA": 1, "compute": 1}
    if has_cupy:
        cuda_options["resources"]["ndarray"] = 1
    cuda_base_spec = {
        "cls": dd.Nanny,
        "options": cuda_options
    }

    def _get_tracing_setup(service_name: str, service_id: str) -> str:
        return (
            f"from libertem.common.tracing import maybe_setup_tracing; "
            f"maybe_setup_tracing(service_name='{service_name}', service_id='{service_id}')"
        )

    if isinstance(cpus, int):
        cpus = tuple(range(cpus))

    for cpu in cpus:
        worker_name = f'{name}-cpu-{cpu}'
        cpu_spec = deepcopy(cpu_base_spec)
        cpu_spec['options']['preload'] = preload + (
            'from libertem.executor.dask import worker_setup; '
            + f'worker_setup(resource="CPU", device={cpu})',
            _get_tracing_setup(worker_name, str(cpu)),
            'libertem.preload',
        )
        workers_spec[worker_name] = cpu_spec

    for service in range(num_service):
        worker_name = f'{name}-service-{service}'
        service_spec = deepcopy(service_base_spec)
        service_spec['options']['preload'] = preload + (
            _get_tracing_setup(worker_name, str(service)),
            'libertem.preload',
        )
        workers_spec[worker_name] = service_spec

    cudas = assign_cudas(cudas)

    for cuda in cudas:
        worker_name = f'{name}-cuda-{cuda}'
        if worker_name in workers_spec:
            num_with_name = sum(n.startswith(worker_name) for n in workers_spec)
            worker_name = f'{worker_name}-{num_with_name - 1}'
        cuda_spec = deepcopy(cuda_base_spec)
        cuda_spec['options']['preload'] = preload + (
            'from libertem.executor.dask import worker_setup; '
            + f'worker_setup(resource="CUDA", device={cuda})',
            _get_tracing_setup(worker_name, str(cuda)),
            'libertem.preload',
        )
        workers_spec[worker_name] = cuda_spec

    return workers_spec


def _run_task(task, params, task_id, threaded_executor, comms_topic: Optional[str]):
    """
    Very simple wrapper function. As dask internally caches functions that are
    submitted to the cluster in various ways, we need to make sure to
    consistently use the same function, and not build one on the fly.

    Without this function, UDFTask->UDF->UDFData ends up in the
    cache, which blows up memory usage over time.
    """
    worker_context = DaskWorkerContext(comms_topic)
    env = Environment(threads_per_worker=1,
                      threaded_executor=threaded_executor,
                      worker_context=worker_context)
    task_result = task(env=env, params=params)
    return {
        "task_result": task_result,
        "task_id": task_id,
    }


def _simple_run_task(task):
    return task()


class CommonDaskMixin:
    client: dd.Client

    def _task_idx_to_workers(self, workers: WorkerSet, idx: int) -> WorkerSet:
        hosts = list(sorted(workers.hosts()))
        host_idx = idx % len(hosts)
        host = hosts[host_idx]
        return workers.filter(lambda w: w.host == host)

    def _future_for_location(
        self,
        task: TaskProtocol,
        locations: WorkerSet,
        resources,
        workers: WorkerSet,
        task_args=None,
        wrap_fn=_run_task,
    ):
        """
        Submit tasks and return the resulting futures

        Parameters
        ----------
        task:
            callable
        locations:
            potential locations to run the task
        resources:
            required resources for the task
        workers : WorkerSet
            Available workers in the cluster
        """
        submit_kwargs = {}
        if task_args is None:
            task_args = {}
        if locations is not None:
            if len(locations) == 0:
                raise ValueError("no workers found for task")
            location_names = locations.names()
        else:
            location_names = None
        submit_kwargs.update({
            'resources': self._validate_resources(workers, resources),
            'workers': location_names,
            'pure': False,
        })
        return self.client.submit(
            wrap_fn, task, *task_args, **submit_kwargs
        )

    def _validate_resources(self, workers, resources):
        # This is set in the constructor of DaskJobExecutor
        if self.lt_resources:
            if not self._resources_available(workers, resources):
                raise ResourceError("Requested resources not available in cluster:", resources)
            result = resources
        else:
            if 'CUDA' in resources:
                raise ResourceError(
                    "Requesting CUDA resource on a cluster without resource management."
                )
            result = {}
        return result

    def _resources_available(self, workers, resources):
        def filter_fn(worker):
            return all(worker.resources.get(key, 0) >= resources[key] for key in resources.keys())

        return len(workers.filter(filter_fn))

    def has_libertem_resources(self):
        workers = self.get_available_workers()

        def has_resources(worker):
            r = worker.resources
            return 'compute' in r and (('CPU' in r and 'ndarray' in r) or 'CUDA' in r)

        return len(workers.filter(has_resources)) > 0

    def _get_future(
        self,
        task: TaskProtocol,
        workers: WorkerSet,
        idx: int,
        params_handle,
        threaded_executor,
        comms_topic: Optional[str]
    ):
        if len(workers) == 0:
            raise RuntimeError("no workers available!")
        params_fut = self._scatter_map[params_handle]
        return self._future_for_location(
            task=task,
            locations=task.get_locations() or self._task_idx_to_workers(
                workers, idx
            ),
            resources=task.get_resources(),
            workers=workers,
            task_args=(
                params_fut,
                idx,
                threaded_executor,
                comms_topic,
            )
        )

    def get_available_workers(self) -> WorkerSet:
        info = self.client.scheduler_info()
        return WorkerSet([
            Worker(
                name=worker['name'],
                host=worker['host'],
                resources=worker['resources'],
                nthreads=worker['nthreads'],
            )
            for worker in info['workers'].values()
        ])

    def get_resource_details(self):
        workers = self.get_available_workers()
        details = {}

        for worker in workers:
            host_name = worker.host
            if worker.name.startswith("tcp"):
                # for handling `dask-worker`
                # `dask-worker` name starts with "tcp"
                #  only supports CPU
                resource = 'cpu'
            else:
                # for handling `libertem-worker`
                r = worker.resources
                if "CPU" in r:
                    resource = 'cpu'
                elif "CUDA" in r:
                    resource = 'cuda'
                else:
                    resource = 'service'

            if host_name not in details.keys():
                details[host_name] = {
                                 'host': host_name,
                                 'cpu': 0,
                                 'cuda': 0,
                                 'service': 0,
                            }
            details[host_name][resource] += 1

        details_sorted = []
        for host in sorted(details.keys()):
            details_sorted.append(details[host])

        return details_sorted


def _dispatch_messages(subscribers: dict[str, list[Callable]], dask_message: tuple[float, dict]):
    """
    Unpacks the Dask message format and forwards the message
    to all subscribed callbacks for that topic (if any)
    """
    timestamp, message = dask_message
    true_topic = message.pop('topic')
    for handler in subscribers.get(true_topic, []):
        handler(true_topic, message)


class DaskJobExecutor(CommonDaskMixin, BaseJobExecutor):
    '''
    Default LiberTEM executor that uses `Dask futures
    <https://docs.dask.org/en/stable/futures.html>`_.

    Parameters
    ----------

    client : distributed.Client
    is_local : bool
        Close the Client and cluster when the executor is closed.
    lt_resources : bool
        Specify if the cluster has LiberTEM resource tags and environment
        variables for GPU processing. Autodetected by default.
    '''
    def __init__(self, client: dd.Client, is_local: bool = False,
                lt_resources: bool = None):
        self.is_local = is_local
        self.client = client
        if lt_resources is None:
            lt_resources = self.has_libertem_resources()
        self.lt_resources = lt_resources
        self._futures = {}
        self._scatter_map = {}

    @contextlib.contextmanager
    def scatter(self, obj):
        # an additional layer of indirection, because we want to be able to
        # redirect keys to new values
        handle = str(uuid.uuid4())
        try:
            fut = self.client.scatter(obj, broadcast=True, hash=False)
            self._scatter_map[handle] = fut
            yield handle
        finally:
            if handle in self._scatter_map:
                del self._scatter_map[handle]

    def scatter_update(self, handle, obj):
        fut = self.client.scatter(obj, broadcast=True, hash=False)
        self._scatter_map[handle] = fut

    def scatter_update_patch(self, handle, patch):
        fut = self._scatter_map[handle]

        def _do_patch(obj):
            if not hasattr(obj, 'patch'):
                raise TypeError(f'object is not patcheable: {obj}')
            obj.patch(patch)

        # can't `client.run` here, as that doesn't resolve the scatter future
        futures = []
        for worker in self.get_available_workers().names():
            futures.append(self.client.submit(_do_patch, fut, pure=False, workers=[worker]))
        for res in dd.as_completed(futures, loop=self.client.loop):
            pass

    def run_tasks(
        self,
        tasks: Iterable[TaskProtocol],
        params_handle: Any,
        cancel_id: Any,
        task_comm_handler: TaskCommHandler,
    ):
        tasks = list(tasks)
        tasks_w_index = list(enumerate(tasks))

        def _id_to_task(task_id):
            return tasks[task_id]

        workers = self.get_available_workers()
        threaded_executor = workers.has_threaded_workers()

        self._futures[cancel_id] = []
        initial = []

        try:
            topic_id = f'topic-{cancel_id}'
            # Wrap all subscriptions into single unique topic
            self.client.subscribe_topic(
                topic_id,
                functools.partial(_dispatch_messages,
                                  task_comm_handler.subscriptions)
            )
        except AttributeError:
            # Dask version does not support structured logs
            # Fall back to partition-level progress updates only
            topic_id = None

        for w in range(int(len(workers))):
            if not tasks_w_index:
                break
            idx, wrapped_task = tasks_w_index.pop(0)
            future = self._get_future(wrapped_task, workers, idx, params_handle,
                                      threaded_executor, topic_id)
            initial.append(future)
            self._futures[cancel_id].append(future)

        try:
            as_completed = dd.as_completed(initial, with_results=True, loop=self.client.loop)
            for future, result_wrap in as_completed:
                if future.cancelled():
                    log.debug(
                        "future %r is cancelled, stopping",
                        future,
                    )
                    del self._futures[cancel_id]
                    raise JobCancelledError()
                result = result_wrap['task_result']
                task = _id_to_task(result_wrap['task_id'])
                if tasks_w_index:
                    idx, wrapped_task = tasks_w_index.pop(0)
                    future = self._get_future(
                        wrapped_task, workers, idx, params_handle, threaded_executor, topic_id
                    )
                    as_completed.add(future)
                    self._futures[cancel_id].append(future)
                yield result, task
        finally:
            if cancel_id in self._futures:
                del self._futures[cancel_id]
            if topic_id is not None:
                self.client.unsubscribe_topic(topic_id)

    def cancel(self, cancel_id):
        log.debug("cancelling with cancel_id=`%s`", cancel_id)
        if cancel_id in self._futures:
            futures = self._futures[cancel_id]
            self.client.cancel(futures)

    def run_each_partition(self, partitions, fn, all_nodes=False):
        """
        Run `fn` for all partitions. Yields results in order of completion.

        Parameters
        ----------

        partitions : List[Partition]
            List of relevant partitions.

        fn : callable
            Function to call, will get the partition as first and only argument.

        all_nodes : bool
            If all_nodes is True, run the function on all nodes that have this partition,
            otherwise run on any node that has the partition. If a partition has no location,
            the function will not be run for that partition if `all_nodes` is True, otherwise
            it will be run on any node.
        """
        def _make_items_all():
            for p in partitions:
                locs = p.get_locations()
                if locs is None:
                    continue
                for workers in locs.group_by_host():
                    yield (lambda: fn(p), workers)

        if all_nodes:
            items = _make_items_all()
        else:
            # TODO check if we should request a compute resource
            items = ((lambda: fn(p), p.get_locations(), {})
                     for p in partitions)
        workers = self.get_available_workers()
        futures = [
            self._future_for_location(*item, workers, wrap_fn=_simple_run_task)
            for item in items
        ]
        # TODO: do we need cancellation and all that good stuff?
        for future, result in dd.as_completed(futures, with_results=True, loop=self.client.loop):
            if future.cancelled():
                raise JobCancelledError()
            yield result

    def run_function(self, fn, *args, **kwargs):
        """
        run a callable :code:`fn` on any worker
        """
        fn_with_args = functools.partial(fn, *args, **kwargs)
        future = self.client.submit(fn_with_args, priority=1, pure=False)
        return future.result()

    def map(self, fn, iterable):
        """
        Run a callable :code:`fn` for each element in :code:`iterable`, on arbitrary worker nodes.

        Parameters
        ----------

        fn : callable
            Function to call. Should accept exactly one parameter.

        iterable : Iterable
            Which elements to call the function on.
        """
        return [future.result()
                for future in self.client.map(fn, iterable, pure=False)]

    def run_each_host(self, fn, *args, **kwargs):
        """
        Run a callable :code:`fn` once on each host, gathering all results into
        a dict host -> result
        """
        # TODO: any cancellation/errors to handle?
        available_workers = self.get_available_workers()

        future_map = {}
        for worker_set in available_workers.group_by_host():
            future_map[worker_set.example().host] = self.client.submit(
                functools.partial(fn, *args, **kwargs),
                priority=1,
                workers=worker_set.names(),
                # NOTE: need pure=False, otherwise the functions will all map to the same
                # scheduler key and will only run once
                pure=False,
            )
        result_map = {
            host: future.result()
            for host, future in future_map.items()
        }
        return result_map

    def run_each_worker(self, fn, *args, **kwargs):
        # Client.run() creates issues on Windows and OS X with Python 3.6
        # FIXME workaround may not be needed anymore for Python 3.7+
        available_workers = self.get_available_workers()

        future_map = {}
        for worker in available_workers:
            future_map[worker.name] = self.client.submit(
                functools.partial(fn, *args, **kwargs),
                priority=1,
                workers=[worker.name],
                # NOTE: need pure=False, otherwise the functions will all map to the same
                # scheduler key and will only run once
                pure=False,
            )
        result_map = {
            name: future.result()
            for name, future in future_map.items()
        }
        return result_map

    def close(self):
        if self.is_local:
            # Client.close won't close the Cluster itself because
            # we provided an external dd.SpecCluster
            self.client.close()
            # Manually close the cluster if not yet torn down
            # use getattr just in case cluster is already gone
            if getattr(self.client, 'cluster', None) is not None:
                self.client.cluster.close(timeout=30)
        # NOTE: distributed already registers atexit handlers for
        # both clients and clusters, this is here to allow manual closure
        # followed by creation of a new Executor without accumulating clusters

    @classmethod
    def connect(cls, scheduler_uri, *args, client_kwargs: Optional[dict] = None, **kwargs):
        """
        Connect to a remote dask scheduler.

        Parameters
        ----------
        scheduler_uri: str
            Compatible with the :code:`address` parameter of :class:`distributed.Client`.
        client_kwargs: dict or None
            Passed as kwargs to :class:`distributed.Client`.
            :code:`client_kwargs['set_as_default']` is set to :code:`False`
            unless specified otherwise to avoid interference with Dask-based workflows.
            Pass :code:`client_kwargs={'set_as_default': True}` to set the Client as the
            default Dask scheduler and keep it running when the Context closes.
        *args, **kwargs: Passed to :class:`DaskJobExecutor`.

        Returns
        -------
        DaskJobExecutor
            the connected JobExecutor
        """
        if client_kwargs is None:
            client_kwargs = {}
        if client_kwargs.get('set_as_default') is None:
            client_kwargs['set_as_default'] = False
        is_local = not client_kwargs['set_as_default']
        client = dd.Client(address=scheduler_uri, **client_kwargs)
        return cls(client=client, is_local=is_local, *args, **kwargs)

    @classmethod
    def make_local(cls, spec: Optional[dict] = None, cluster_kwargs: Optional[dict] = None,
            client_kwargs: Optional[dict] = None, preload: Optional[tuple[str]] = None):
        """
        Spin up a local dask cluster

        Parameters
        ----------
        spec
            Dask cluster spec, see
            https://distributed.dask.org/en/stable/api.html#distributed.SpecCluster
            for more info.
            :func:`libertem.utils.devices.detect` allows to detect devices that can be used
            with LiberTEM, and :func:`cluster_spec` can be used to create a :code:`spec`
            with customized parameters.
        cluster_kwargs
            Passed to :class:`distributed.SpecCluster`.
        client_kwargs
            Passed to :class:`distributed.Client`. Pass
            :code:`client_kwargs={'set_as_default': True}` to set the Client as the
            default Dask scheduler.
        preload: Optional[Tuple[str]]
            Passed to :func:`cluster_spec` if :code:`spec` is :code:`None`.

        Returns
        -------
        DaskJobExecutor
            the connected JobExecutor
        """

        # Distributed doesn't adjust the event loop policy when being run
        # from within pytest as of version 2.21.0. For that reason we
        # adjust the policy ourselves here.
        adjust_event_loop_policy()

        if spec is None:
            from libertem.utils.devices import detect
            spec = cluster_spec(**detect(), preload=preload)
        else:
            if preload is not None:
                raise ValueError(
                    "Passing both spec and preload is not supported. "
                    "Instead, include preloading specification in the spec"
                )
        if client_kwargs is None:
            client_kwargs = {}
        if client_kwargs.get('set_as_default') is None:
            client_kwargs['set_as_default'] = False

        if cluster_kwargs is None:
            cluster_kwargs = {}
        if cluster_kwargs.get('silence_logs') is None:
            cluster_kwargs['silence_logs'] = logging.WARN

        dist_log_level = dask.config.get('distributed.logging.distributed', default=None)
        if dist_log_level is None:
            dist_log_level = cluster_kwargs['silence_logs']

        with set_num_threads_env(n=1), set_worker_log_level(dist_log_level):
            # Mitigation for https://github.com/dask/distributed/issues/6776
            with dask.config.set({"distributed.worker.profile.enabled": False}):
                cluster = dd.SpecCluster(workers=spec, **(cluster_kwargs or {}))
                client = dd.Client(cluster, **(client_kwargs or {}))
                client.wait_for_workers(len(spec))

        is_local = not client_kwargs['set_as_default']

        return cls(client=client, is_local=is_local, lt_resources=True)

    def __enter__(self):
        return self

    def __exit__(self, exc_type, exc_val, exc_tb):
        self.close()


class AsyncDaskJobExecutor(AsyncAdapter):
    def __init__(self, wrapped=None, *args, **kwargs):
        if wrapped is None:
            wrapped = DaskJobExecutor(*args, **kwargs)
        super().__init__(wrapped)

    @classmethod
    async def connect(cls, scheduler_uri, *args, **kwargs):
        executor = await sync_to_async(functools.partial(
            DaskJobExecutor.connect,
            scheduler_uri=scheduler_uri,
            *args,
            **kwargs,
        ))
        return cls(wrapped=executor)

    @classmethod
    async def make_local(cls, spec=None, cluster_kwargs=None, client_kwargs=None):
        executor = await sync_to_async(functools.partial(
            DaskJobExecutor.make_local,
            spec=spec,
            cluster_kwargs=cluster_kwargs,
            client_kwargs=client_kwargs,
        ))
        return cls(wrapped=executor)


def cli_worker(
    scheduler,
    local_directory,
    cpus,
    cudas,
    has_cupy,
    name,
    log_level,
    preload: tuple[str, ...]
):
    import asyncio

    options = {
        "silence_logs": log_level,
        "local_directory": local_directory
    }

    spec = cluster_spec(
        cpus=cpus, cudas=cudas, has_cupy=has_cupy, name=name,
        options=options, preload=preload,
    )

    async def run(spec):
        # Mitigation for https://github.com/dask/distributed/issues/6776
        with dask.config.set({"distributed.worker.profile.enabled": False}):
            workers = []
            for name, spec in spec.items():
                cls = spec['cls']
                workers.append(
                    cls(scheduler, name=name, **spec['options'])
                )
            import asyncio
            await asyncio.gather(*workers)
            for w in workers:
                await w.finished()

    asyncio.get_event_loop().run_until_complete(run(spec))
