Setting up, configuring and deploying Dask Distributed

 1def load_configuration():
 2    import os
 3    import yaml
 4
 5    from dask.config import collect_env, config, update
 6
 7    # Load config overrides from file
 8    config_path = os.path.join(os.path.dirname(__file__), 'dask_configuration.yaml')
 9    with open(config_path) as f:
10        overrides = yaml.safe_load(f)
11        update(config, overrides)
12
13    # Override from environment variables
14    collect_env()

distributed:
  client:
    heartbeat: "10s"
    scheduler-info-interval: "5s"
  comm:
    timeouts:
      connect: "30s"
      tcp: "45s"
    retry:
      count: 3
      delay:
        min: "2s"
        max: "30s"
    socket-backlog: 4096
  deploy:
    lost-worker-timeout: "35s"
  scheduler:
    events-cleanup-delay: "1800s"
    locks:
      lease-validation-interval: "3s"
      lease-timeout: "1260s"
    http:
      routes: [
        'distributed.http.scheduler.prometheus',
        'distributed.http.scheduler.info',
        'distributed.http.scheduler.json',
        'distributed.http.health',
        'distributed.http.statics',
        'my_project.dask_routes.scheduler_routes',
      ]
    preload: [
      'my_project.dask_preload',
    ]
  worker:
    connections:
      outgoing: 50
      incoming: 25
    memory:
      target: 0.75
      spill: 0.85
      pause: 0.90
      terminate: 0.95
    http:
      routes: [
        'distributed.http.worker.prometheus',
        'distributed.http.health',
        'distributed.http.statics',
        'my_project.dask_routes.worker_routes',
      ]
    preload: [
      'my_project.dask_preload',
    ]
  admin:
    log-length: 3000
    log-format: '%(asctime)s - %(name)s - %(levelname)s - %(message)s'

import logging
from types import TracebackType
from typing import Optional, cast

from dask.utils import funcname
from distributed import Worker
from distributed.diagnostics.plugin import WorkerPlugin


class WorkerMonitor(WorkerPlugin):
    def __init__(self):
        super().__init__()
        self.logger = logging.getLogger('dask.worker.worker_monitor')

    def setup(self, worker: Worker):
        self.worker = worker  # noqa

    def transition(self, key: str, start: str, finish: str, *args, **kwargs):
        if finish == 'error' and key in self.worker.exceptions and key in self.worker.tasks:
            serialized_exc = self.worker.exceptions[key]

            if hasattr(serialized_exc, 'data') and isinstance(serialized_exc.data, Exception):
                tb = self.worker.tracebacks.get(key)
                exc = (
                    serialized_exc.data.with_traceback(tb.data)
                    if (tb and hasattr(tb, 'data') and isinstance(tb.data, TracebackType))
                    else serialized_exc.data
                )
                exc = cast(Exception, exc)

                function, fn_args, fn_kwargs = self.worker.tasks[key]
                function_name = str(funcname(function))

                extra_data = {
                    "task_id": key,
                    "task_name": function_name,
                    "args": fn_args,
                    "kwargs": fn_kwargs,
                    'dask_worker_name': self.worker.name,
                }

                # TODO: your exception reporting and tracing
                my_exception_reporter.report(exc, extra_data)

from typing import Union

import click
from distributed import Client, Scheduler, Worker


@click.command()
def dask_setup(service: Union[Scheduler, Worker]):
    """
    References:
    * Lifecycle and CLI integration: https://docs.dask.org/en/latest/setup/custom-startup.html
    * Developing plugins: https://distributed.dask.org/en/latest/plugins.html
    """

    # Scheduler specific
    if isinstance(service, Scheduler):
        # Graceful Scheduler shutdown
        # (remove once https://github.com/dask/distributed/pull/3332 gets merged & released)
        async def on_signal(_signum):
            await service.close()

        # use the reference implementation from `distributed`
        install_signal_handlers(service.loop, cleanup=on_signal)

    # Worker specific
    elif isinstance(service, Worker):
        # Plugin for monitoring & metrics
        worker_monitor_plugin = WorkerMonitor()
        with Client(
            address=service.scheduler.address,
            timeout=30,
            name='worker-plugin-setup-client',
        ) as client:
            client.register_worker_plugin(worker_monitor_plugin, name='worker-monitor')

from functools import partial
from typing import Any, Dict, List, Optional

from distributed import Scheduler
from distributed.http.proxy import GlobalProxyHandler
from distributed.scheduler import WorkerState
from tornado.web import RequestHandler


class WorkerOnlyProxyHandler(GlobalProxyHandler):
    """
    Endpoint that allows proxying Workers from the Scheduler.
    Useful to remotely shutdown a specific worker or access worker's metrics, logs and dashboards.
    """

    def initialize(self, dask_server: Optional[Scheduler] = None, extra: Optional[Dict[Any, Any]] = None):
        super().initialize(dask_server=dask_server, extra=extra)

        worker_hosts: List[str] = (
            [
                worker_state.host for worker_state in dask_server.workers.values()  # type: WorkerState
            ]
            if dask_server
            else []
        )

        # override the whitelist function to proxy only to workers
        self.host_whitelist = partial(self.whitelist_workers, worker_hosts=worker_hosts)

    @staticmethod
    def whitelist_workers(_handler: RequestHandler, host: str, *, worker_hosts: List[str]):
        return any(host in worker_host for worker_host in worker_hosts)


# Export the HTTP routes
#
# https://docs.dask.org/en/latest/configuration-reference.html#distributed.scheduler.http.routes
# https://distributed.dask.org/en/latest/http_services.html
routes = [(r"proxy/(\d+)/(.*?)/(.*)", WorkerOnlyProxyHandler, {})]

import logging

from distributed import Worker
from distributed.http.utils import RequestHandler

logger = logging.getLogger("distributed.request_handler")


class TerminationHandler(RequestHandler):
    """
    Custom HTTP handler to trigger a graceful shutdown via Kubernetes's Lifecycle Hooks
    reference: https://kubernetes.io/docs/concepts/containers/container-lifecycle-hooks/
    """

    def get(self):
        self.server: Worker  # add typing information (Worker extends distributed.core.Server)

        logger.info(f"Lifecycle hook triggered. Initiating graceful shutdown of {self.server.name}.")

        self.server.io_loop.add_callback(self.server.close_gracefully)

        self.write({'message': 'Shutting down...', 'extra': {'worker_name': self.server.name}})
        self.set_header("Content-Type", "application/json")


# Export the HTTP routes
#
# https://docs.dask.org/en/latest/configuration-reference.html#distributed.worker.http.routes
# https://distributed.dask.org/en/latest/http_services.html
routes = [('graceful-shutdown', TerminationHandler, {})]