3-distributed-scheduler.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "9d19f489-59e6-48c7-aade-25530dcf253c",
   "metadata": {},
   "source": [
    "<img src=\"images/dask_horizontal.svg\" width=\"45%\" alt=\"Dask logo\\\">\n",
    "\n",
    "# Dask clusters\n",
    "\n",
    "This notebook covers Dask's distributed clusters in more detail. We provide a more in depth look at the components of a cluster, illustrate how to inspect the internal state of a cluster, and how you can extend the functionality of your cluster using Dask's plugin system."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c45762cb-a5a2-4f7c-ae29-02f05c78ba89",
   "metadata": {},
   "source": [
    "# Cluster overview\n",
    "\n",
    "In this section we'll discuss:\n",
    "\n",
    "1. The different components which make up a Dask cluster\n",
    "2. Survey different ways to launch a cluster\n",
    "\n",
    "## Components of a cluster\n",
    "\n",
    "A Dask cluster is composed of three different types of objects:\n",
    "\n",
    "1. **Scheduler**: A single, centralized scheduler process which responds to requests for computations, maintains relavant state about tasks and worker, and sends tasks to workers to be computed.\n",
    "2. **Workers**: One or more worker processes which compute tasks and store/serve their results.\n",
    "3. **Clients**: One or more client objects which are the user-facing entry point to interact with the cluster.\n",
    "\n",
    "<img src=\"images/distributed-overview.svg\"\n",
    "     width=\"90%\"\n",
    "     alt=\"Dask components\\\">\n",
    "\n",
    "A couple of notes about workers:\n",
    "\n",
    "- Each worker runs in its own Python process. Each worker Python process has its own `concurrent.futures.ThreadPoolExecutor` which is uses to compute tasks in parallel. The same threads vs. processes considerations we discussed earlier also apply to Dask workers.\n",
    "- There's actually a fourth cluster object which is often not discussed: the **Nanny**. By default Dask workers are launched and managed by a separate nanny process. This separate process allows workers to restart themselves if you want to use the `Client.restart` method, or to restart workers automatically if they get above a certain memory limit threshold.\n",
    "\n",
    "#### Related Documentation\n",
    "\n",
    "- [Cluster architecture](https://distributed.dask.org/en/latest/#architecture)\n",
    "- [Journey of a task](https://distributed.dask.org/en/latest/journey.html)\n",
    "\n",
    "## Deploying Dask clusters\n",
    "\n",
    "Deploying a Dask cluster means launching scheduler, worker, and client processes and setting up the appropriate network connections so these processes can communicate with one another. Dask clusters can be lauched in a few different ways which we highlight in the following sections.\n",
    "\n",
    "\n",
    "### Manual setup\n",
    "\n",
    "Launch a scheduler process using the `dask-scheduler` command line utility:\n",
    "\n",
    "```terminal\n",
    "$ dask-scheduler\n",
    "Scheduler at:   tcp://192.0.0.100:8786\n",
    "```\n",
    "\n",
    "and then launch several workers by using the `dask-worker` command and providing them the address of the scheduler they should connect to:\n",
    "\n",
    "```terminal\n",
    "$ dask-worker tcp://192.0.0.100:8786\n",
    "Start worker at:  tcp://192.0.0.1:12345\n",
    "Registered to:    tcp://192.0.0.100:8786\n",
    "\n",
    "$ dask-worker tcp://192.0.0.100:8786\n",
    "Start worker at:  tcp://192.0.0.2:40483\n",
    "Registered to:    tcp://192.0.0.100:8786\n",
    "\n",
    "$ dask-worker tcp://192.0.0.100:8786\n",
    "Start worker at:  tcp://192.0.0.3:27372\n",
    "Registered to:    tcp://192.0.0.100:8786\n",
    "```\n",
    "\n",
    "### Python API (advanced)\n",
    "\n",
    "⚠️ **Warning**: Creating `Scheduler` / `Worker` objects explicitly in Python is rarely needed in practice and is intended for more advanced users ⚠️"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9adfaf00-139c-4a4e-8971-318c9d8c31c7",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "from dask.distributed import Scheduler, Worker, Client"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "392a2886-05e0-4af7-8cb2-63015ba13e5a",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "# Launch a scheduler\n",
    "async with Scheduler() as scheduler: # Launch a scheduler\n",
    "    # Launch a worker which connects to the scheduler\n",
    "    async with Worker(scheduler.address) as worker:\n",
    "        # Launch a client which connects to the scheduler\n",
    "        async with Client(scheduler.address, asynchronous=True) as client:\n",
    "            result = await client.submit(sum, range(100))\n",
    "            print(f\"{result = }\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ec8031ac-58a0-40bb-a0b1-07847d54b964",
   "metadata": {},
   "source": [
    "### Cluster managers (recommended)\n",
    "\n",
    "Dask has the notion of cluster manager objects. Cluster managers offer a consistent interface for common activities like adding/removing workers to a cluster, retrieving logs, etc."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "304c076c-5f10-4760-88b2-cf4dff1aed99",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "from dask.distributed import LocalCluster"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "584a2ae1-32ad-4ca0-b563-c5f58595c250",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "# Launch a scheduler and 4 workers on my local machine\n",
    "cluster = LocalCluster(n_workers=4, threads_per_worker=2)\n",
    "cluster"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "41daa9b9-51c5-4996-85a5-144f19714389",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "# Scale up to 10 workers\n",
    "cluster.scale(10)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c8937b5b-2559-4eb1-96ae-d76b385526c2",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "# Scale down to 2 workers\n",
    "cluster.scale(2)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3b6098c3-1f8b-4e42-8d0a-996f2b44b1e3",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "# Retrieve cluster logs\n",
    "cluster.get_logs()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e350ead3-229a-4974-bd89-7174d99edbef",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "# Shut down cluster\n",
    "cluster.close()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "94169465-79c0-4063-b7bf-dbc9264cdfe4",
   "metadata": {},
   "source": [
    "There are several projects in the Dask ecosystem for easily deploying clusters on commonly used computing resources:\n",
    "\n",
    "- [Dask-Kubernetes](https://kubernetes.dask.org/en/latest/) for deploying Dask using native Kubernetes APIs\n",
    "- [Dask-Cloudprovider](https://cloudprovider.dask.org/en/latest/) for deploying Dask clusters on various cloud platforms (e.g. AWS, GCP, Azure, etc.)\n",
    "- [Dask-Yarn](https://yarn.dask.org/en/latest/) for deploying Dask on YARN clusters\n",
    "- [Dask-MPI](http://mpi.dask.org/en/latest/) for deploying Dask on existing MPI environments\n",
    "- [Dask-Jobqueue](https://jobqueue.dask.org/en/latest/) for deploying Dask on job queuing systems (e.g. PBS, Slurm, etc.)\n",
    "\n",
    "Launching clusters with any of these projects follows a similar pattern as using Dask's built-in `LocalCluster`:\n",
    "\n",
    "```python\n",
    "# Launch a Dask cluster on a Kubernetes cluster\n",
    "from dask_kubernetes import KubeCluster\n",
    "cluster = KubeCluster(...)\n",
    "\n",
    "# Launch a Dask cluster on AWS Fargate\n",
    "from dask_cloudprovider.aws import FargateCluster\n",
    "cluster = FargateCluster(...)\n",
    "\n",
    "# Launch a Dask cluster on a PBS job queueing system\n",
    "from dask_jobqueue import PBSCluster\n",
    "cluster = PBSCluster(...)\n",
    "```\n",
    "\n",
    "Additionally, there are companies like [Coiled](https://coiled.io) and [Saturn Cloud](https://www.saturncloud.io) which have Dask deployment-as-a-service offerings. *Disclaimer*: Some of the instructors for this tutorial are employed by Coiled. \n",
    "\n",
    "#### Related Documentation\n",
    "\n",
    "- [Deploy Dask Clusters](https://docs.dask.org/en/stable/deploying.html)\n",
    "- [Cluster setup](https://docs.dask.org/en/latest/setup.html)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "20ba2911-0ac4-43bd-bdb9-54a2806e5ea9",
   "metadata": {},
   "source": [
    "# Inspecting a cluster's state\n",
    "\n",
    "In this section we'll:\n",
    "\n",
    "1. Familiarize ourselves with Dask's scheduler and worker processes\n",
    "2. Explore the various state that's tracked throughout the cluster\n",
    "3. Learn how to inspect remote scheduler and worker processes\n",
    "\n",
    "Dask has a a variety of ways to provide users insight into what's going on during their computations. For example, Dask's [diagnositc dashboard](https://docs.dask.org/en/latest/diagnostics-distributed.html) displays real-time information about what tasks are current running, overal progress on a computation, worker CPU and memory load, statistical profiling information, and much more. Additionally, Dask's [performance reports](https://distributed.dask.org/en/latest/diagnosing-performance.html#performance-reports) allow you to save the diagnostic dashboards as static HTML plots. Performance reports are particularly useful when benchmarking/profiling workloads or when sharing workload performance with colleagues."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ee7a1ae0-f029-41a8-90bc-7d6a5e6651e8",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "from dask.distributed import LocalCluster, Client, Worker\n",
    "\n",
    "cluster = LocalCluster(worker_class=Worker, processes=True) \n",
    "client = Client(cluster)\n",
    "client"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "33ce7759-2113-4dcc-9581-2a473eead63b",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "import dask.array as da\n",
    "from dask.distributed import performance_report\n",
    "\n",
    "with performance_report(\"my_report.html\"):\n",
    "    x = da.random.random((10_000, 10_000), chunks=(1_000, 1_000))\n",
    "    result = (x + x.T).mean(axis=0).mean()\n",
    "    result.compute()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "bed8743b-a34e-488a-bc51-abe923914cf0",
   "metadata": {},
   "source": [
    "These are invaluable tools and we highly recommend utilizing them. Often times Dask's dashboard is totally sufficient to understand the performance of your computations.\n",
    "\n",
    "However, sometimes it can be useful to dive more deeply into the internals of your cluster and directly inspect the state of your scheduler and workers. Let's start by submitting some tasks to the cluster to be computed."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "98cd7868-ceb2-4066-bb8f-094ad3fc0053",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "import random\n",
    "\n",
    "def double(x):\n",
    "    random.seed(x)\n",
    "    # Simulate some random task failures\n",
    "    if random.random() < 0.1:\n",
    "        raise ValueError(\"Oh no!\")\n",
    "    return 2 * x\n",
    "\n",
    "futures = client.map(double, range(50))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3b8a3c19-97d2-47a9-8216-3f73abe6189a",
   "metadata": {},
   "source": [
    "One of the nice things about `LocalCluster` is it gives us direct access the `Scheduler` Python object. This allows us to easily inspect the scheduler directly."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5b842c4f-64c6-4a0b-b068-45332552d6a9",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "scheduler = cluster.scheduler\n",
    "scheduler"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6b8e8155-68a2-4cfc-8cfb-424180e98bdb",
   "metadata": {},
   "source": [
    "ℹ️ Note that often times you won't have direct access to the `Scheduler` Python object (e.g. when the scheduler is running on separate machine). In these cases it's still possible to inspect the scheduler and we will discuss how to do this later on.\n",
    "\n",
    "The scheduler tracks **a lot** of state. Let's start to explore the scheduler to get a sense for what information it keeps track of."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a4a43b6c-210c-41a0-bfff-38c38bfb6316",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "scheduler.address   # Scheduler's address"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "eb25dc76-df8c-418e-84ea-ba3d33fcd373",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "scheduler.time_started   # Time the scheduler was started"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "59c8b1be-3adb-4cdd-893c-f83d71a75a92",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "dict(scheduler.workers)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f9f68f77-69c6-4b17-90ce-cf179af7d287",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "worker_state = next(iter(scheduler.workers.values()))\n",
    "worker_state"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f5ba8180-8ef7-4bc7-ae0a-cf7b9708c9e6",
   "metadata": {},
   "source": [
    "Let's take a look at the `WorkerState` attributes"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a73d206f-7a65-49f0-8361-3f9e37f5a73a",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "worker_state.address   # Worker's address"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c0d58de4-ec72-4d16-a77b-f35a05d74563",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "worker_state.status   # Current status of the worker (e.g. \"running\", \"closed\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9b0add25-1edc-4bad-92dd-c2b18947d451",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "worker_state.nthreads   # Number of threads in the worker's `ThreadPoolExecutor`"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "24be2aab-a7c7-454d-91f9-78897e32eadb",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "worker_state.executing   # Dictionary of all tasks which are currently being processed, along with the current duration of the task"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8e414215-dc91-4468-af2f-4119945bca3a",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "worker_state.metrics   # Various metrics describing the current state of the worker"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e2bbacd5-a368-4566-912c-4a4db4b8a211",
   "metadata": {},
   "source": [
    "Workers check in with the scheduler inform it when certain event occur (e.g. when a worker has completed a task) so the scheduler can update it's internal state."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "abf2af29-be3e-4d78-a5c2-ca595d380357",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "worker_state.last_seen"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "be601b2d-7047-44b2-81c2-5dabd344053c",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "import time\n",
    "\n",
    "for _ in range(10):\n",
    "    print(f\"{worker_state.last_seen = }\")\n",
    "    time.sleep(1)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "13b53603-8d32-40c8-b71d-b3ed1af57848",
   "metadata": {},
   "source": [
    "In addition to the state of each worker, the scheduler also tracks information for each task it has been asked to run."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ca9ddfa1-7ab7-4585-961d-01522c343b42",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "scheduler.tasks"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ac277222-ac14-4743-be25-d2460f5d9d9f",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "task_state = next(iter(scheduler.tasks.values()))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9d8093d0-9542-42d7-a417-1e9b59a6970d",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "task_state"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1bc9a48c-6c3a-4f39-be74-fae5e26ea645",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "task_state.key   # Task's name (unique identifier)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "cfe7c942-172b-4e54-9548-8e94c17da4c0",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "task_state.state   # Task's state (e.g. \"memory\", \"waiting\", \"processing\", \"erred\", etc.)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b9cd9140-19ce-45dd-adc0-d74c7a671ba1",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "task_state.who_has   # Set of workers (`WorkerState`s) who have this task's result in memory"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "fce9b8e3-e8c7-49d0-8343-00d899c35a8a",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "task_state.nbytes   # The number of bytes of the result of this finished task"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4b2ec17e-8bb3-47a7-b078-869b0019f6ae",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "task_state.type   # The type of the the task's result (as a string)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1ee16f2c-3f95-41f6-ad80-8530b1c33802",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "task_state.retries   # The number of times this task can automatically be retried in case of failure"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f498bef9-988d-4991-981f-d62158dec477",
   "metadata": {},
   "source": [
    "## Exercise 1\n",
    "\n",
    "Spend the next 5 minutes continuing to explore the attributes the scheduler keeps track of. Try to answer the following questions:\n",
    "\n",
    "1. What are the keys for the tasks which failed?\n",
    "2. How many tasks successfully ran on each worker?"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0f35a9e7-05e6-4395-9551-bf6ba39d6aad",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "# What are the keys for the tasks which failed?\n",
    "# Your solution goes here"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4b4e6c15-bf52-4bf5-b841-9646770ac99e",
   "metadata": {
    "jupyter": {
     "source_hidden": true
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "# Solution to \"What are the keys for the tasks which failed?\"\n",
    "erred_tasks = [key for key, ts in scheduler.tasks.items() if ts.state == \"erred\"]\n",
    "erred_tasks"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5d22af01-a1e8-400e-b049-9889c24142ae",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "# How many tasks successfull ran on each worker?\n",
    "# Your solution goes here"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b698f481-b9d6-4e9c-ac0c-325d1ef725d8",
   "metadata": {
    "jupyter": {
     "source_hidden": true
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "# Solution to \"How many tasks successfull ran on each worker?\"\n",
    "from collections import defaultdict\n",
    "\n",
    "erred_tasks = [key for key, ts in scheduler.tasks.items() if ts.state == \"erred\"]\n",
    "counter = defaultdict(int)\n",
    "for key, ts in scheduler.tasks.items():\n",
    "    if key in erred_tasks:\n",
    "        continue\n",
    "    for worker in ts.who_has: \n",
    "        counter[worker] += 1\n",
    "print(counter)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "52e806fc-f8ee-431e-879c-6be9b3ea5d9c",
   "metadata": {
    "jupyter": {
     "source_hidden": true
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "# Solution 2 \n",
    "counter = {address: worker_state.metrics['task_counts']['memory'] \n",
    "           for address, worker_state in scheduler.workers.items()}\n",
    "print(counter)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5e07c7c2-f157-4fbb-82cb-81b59fff7da6",
   "metadata": {},
   "source": [
    "In addition to inspecting the scheduler, we can also investigate the state of each of our workers."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "fbc95933-bff8-4a8e-9b28-430c4a042034",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "cluster.workers"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "42800bd2-8ef2-456a-87d7-0354cf7c9626",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "worker = next(iter(cluster.workers.values()))\n",
    "worker"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b3f5ccb3-4495-4290-9209-25b2687c951e",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "worker.address   # Worker's address"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e89df365-8547-4157-823c-0b33c6ef83e9",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "worker.state.executing_count   # Number of tasks the worker is currenting computing"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "969a7bed-a6e3-4708-9f95-1900d48b27f5",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "worker.state.executed_count   # Running total of all tasks processed on this worker"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "aaeebb27-e9c3-456b-9267-c5180559e82e",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "worker.state.nthreads   # Number of threads in the worker's ThreadPoolExecutor"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ab866a75-7e95-4c5f-b3ff-e54306492fa2",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "worker.executor   # Worker's ThreadPoolExecutor where it computes tasks"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "257fd911-8d37-421a-87ad-1f87e4f68f3e",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "worker.keys()   # Keys the worker currently has in memory"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c1b9afab-edad-4a99-8471-71a943f81d35",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "worker.data   # Where the worker stores task results"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5d45c97e-d8c5-4a50-882c-689ec62378b9",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "{key: worker.data[key] for key in worker.keys()} # results stored on the worker"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c1e20d0f-38d7-4694-a65e-e223557d5d3a",
   "metadata": {},
   "source": [
    "## Accessing remote scheduler and workers\n",
    "\n",
    "As we noted earlier, often times you won't have direct access to the `Scheduler` or `Worker` Python objects for your cluster. However, in these cases it's still possible to examine the state of the scheduler and workers in your cluster using the `Client.run` ([docs](https://distributed.dask.org/en/latest/api.html#distributed.Client.run)) and `Client.run_on_scheduler`([docs](https://distributed.dask.org/en/latest/api.html#distributed.Client.run_on_scheduler)) methods.\n",
    "\n",
    "`Client.run` allows you to run a function on worker processes in your cluster. If the function has a `dask_worker` parameter, then that variable will be populated with the `Worker` instance when the function is run. Likewise, `Client.run_on_scheduler` allows you to run a function on the scheduler processes in your cluster. If the function has a `dask_scheduler` parameter, then that variable will be populated with the `Scheduler` instance when the function is run.\n",
    "\n",
    "Let's look at an examples of custom function. If the function has a `dask_worker` parameter ..."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f2567d89-6409-42ec-b147-c24ab84ac643",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "def get_worker_name(dask_worker):\n",
    "    return dask_worker.name\n",
    "\n",
    "client.run(get_worker_name)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "09ac4f6b-c0ae-4c7e-a66b-1f39a3ceb83b",
   "metadata": {},
   "source": [
    "Similarly, we can do the same thing on the scheduler by using `Client.run_on_scheduler`"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0daf0c15-c20d-421e-a21c-3ee27df6e723",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "def get_erred_tasks(dask_scheduler):\n",
    "    return [key for key, ts in dask_scheduler.tasks.items() if ts.state == \"erred\"]\n",
    "\n",
    "client.run_on_scheduler(get_erred_tasks)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "780f7e4f-8ad2-4e93-ac86-e2d8c5331902",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "client.close()\n",
    "cluster.close()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "aec78fc7-3125-4d40-a03e-3595c384382d",
   "metadata": {},
   "source": [
    "#### Related Documentation\n",
    "\n",
    "- [Dask worker](https://distributed.dask.org/en/latest/worker.html)\n",
    "- [Scheduling state](https://distributed.dask.org/en/latest/scheduling-state.html)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "71b05ace-1e10-4716-ae83-f7c659c46d28",
   "metadata": {},
   "source": [
    "# Extending the scheduler and workers: Dask's plugin system\n",
    "\n",
    "In this section we'll siscuss Dask's scheduler and worker plugin systems and write our own plugin to extend the scheduler's functionality.\n",
    "\n",
    "So far we've primarily focused on inspecting the state of a cluster. However, there are times when it's useful to extend the functionality of the scheduler and/or workers in a cluster. To help facilitate this, Dask has scheduler and worker plugin systems which enable you to hook into different events that happen throughout a cluster's lifecycle. This allows you to run custom code when a specific type of event occurs on the cluster.\n",
    "\n",
    "Specifically, the [scheduler plugin system](https://distributed.dask.org/en/latest/plugins.html#scheduler-plugins) enables you run custom code when the following events occur:\n",
    "\n",
    "1. Scheduler starts, stops, or is restarted\n",
    "2. Client connects or disconnects to the scheduler\n",
    "3. Workers enters or leaves the cluster\n",
    "4. When a new task enters the scheduler\n",
    "5. When a task changes state (e.g. from \"processing\" to \"memory\")\n",
    "\n",
    "While the [worker plugin system](https://distributed.dask.org/en/latest/plugins.html#worker-plugins) enables you run custom code when the following events occur:\n",
    "\n",
    "1. Worker starts or stops\n",
    "2. When a worker releases a task\n",
    "3. When a task changes state (e.g. \"processing\" to \"memory\")\n",
    "\n",
    "Implementing your own custom plugin consists of creating a Python class with certain methods (each method corresponds to a particular lifecycle event)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "de6f767a-52e6-4c9f-a26a-7a642b4bb953",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "from distributed import SchedulerPlugin, WorkerPlugin"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6f35b157-5502-4809-a1e6-5991799c8d31",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "# Lifecycle SchedulerPlugin methods\n",
    "[attr for attr in dir(SchedulerPlugin) if not attr.startswith(\"_\")]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "69ac1db2-1170-49dd-a043-a646df6cedc6",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "# Lifecycle WorkerPlugin methods\n",
    "[attr for attr in dir(WorkerPlugin) if not attr.startswith(\"_\")]"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "131294ca-74ee-4754-bd3e-15eba1de758b",
   "metadata": {},
   "source": [
    "For the exact signature of each method, please refer to the [`SchedulerPlugin`](https://distributed.dask.org/en/latest/plugins.html#scheduler-plugins) and [`WorkerPlugin`](https://distributed.dask.org/en/latest/plugins.html#worker-plugins) documentation.\n",
    "\n",
    "Let's looks at an example scheduler plugin."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4fc97c90-a858-40ea-9f31-e7f90aba3e85",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "class Counter(SchedulerPlugin):\n",
    "    \"\"\"Keeps a running count of the total number of completed tasks\"\"\"\n",
    "    def __init__(self):\n",
    "        self.n_tasks = 0\n",
    "\n",
    "    def transition(self, key, start, finish, *args, **kwargs):\n",
    "        if start == \"processing\" and finish == \"memory\":\n",
    "            self.n_tasks += 1\n",
    "\n",
    "    def restart(self, scheduler):\n",
    "        self.n_tasks = 0"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "51ae7356-5a48-4315-abb4-343378b88f1f",
   "metadata": {},
   "source": [
    "Check documentation on [Task States](https://distributed.dask.org/en/stable/scheduling-state.html#scheduler-task-state)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e2586c29-e5e3-495d-84aa-d30f02639ff7",
   "metadata": {},
   "source": [
    "To add a custom scheduler plugin to your cluster, use the `Scheduler.add_plugin` method:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "56d0b325-2524-4c33-809c-5b853c94594b",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "# Create LocalCluster and Client\n",
    "cluster = LocalCluster()\n",
    "client = Client(cluster)\n",
    "\n",
    "# Instantiate and add the Counter to our cluster\n",
    "counter = Counter()\n",
    "cluster.scheduler.add_plugin(counter)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "485d1388-4cec-4a16-a2d1-2f2e5d3d2448",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "counter.n_tasks"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "220a76d3-25d4-4a3b-800e-fea64b055cb7",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "from distributed import wait\n",
    "futures = client.map(lambda x: x + 1, range(27))\n",
    "wait(futures);"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5a92ad14-2976-4552-b654-8a7ad8645d8a",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "counter.n_tasks"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a1184a1e-bd18-4384-abf3-9acb1ffeccde",
   "metadata": {},
   "outputs": [],
   "source": [
    "client.restart()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f91f910d-022f-4e51-b762-36113c4f4184",
   "metadata": {},
   "outputs": [],
   "source": [
    "counter.n_tasks"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c4a05fc9-8292-48c7-a67d-21ac536e4d3c",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "client.close()\n",
    "cluster.close()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2137bfc5-255d-4047-821d-54a10791b451",
   "metadata": {},
   "source": [
    "This is a relatively straightforward plugin one could write. Let's look at some of the `distributed`s built-in worker plugins to see two more real-world example."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "02ac484a-f5a4-4c3b-8af2-80f79d98d749",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "from distributed import Environ, UploadFile, UploadDirectory, PipInstall, CondaInstall, PackageInstall"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1060ed83-4a35-4cb4-8e41-49ce2f0b51a5",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "Environ?? "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "09e73aa0-3e1c-47d6-9372-2f6df644f87f",
   "metadata": {},
   "outputs": [],
   "source": [
    "PackageInstall??"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "530d27a2-8f93-4735-bc9f-b4e508b7ad53",
   "metadata": {},
   "source": [
    "To add a custom worker plugin to your cluster, use the `Client.register_worker_plugin` method."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "efde3b89-47c8-4d75-b015-be7641284ea6",
   "metadata": {},
   "source": [
    "## Exercise 2\n",
    "\n",
    "Over the next 10 minutes, create a `TaskTimerPlugin` scheduler plugin which keeps tracks of how long each task takes to run.\n",
    "\n",
    "```python\n",
    "\n",
    "class TaskTimerPlugin(SchedulerPlugin):\n",
    "    ...\n",
    "\n",
    "# Create LocalCluster and Client\n",
    "cluster = LocalCluster()\n",
    "client = Client(cluster)\n",
    "\n",
    "# Instantiate and add the TaskTimerPlugin to our cluster\n",
    "plugin = TaskTimerPlugin()\n",
    "cluster.scheduler.add_plugin(plugin)\n",
    "\n",
    "import dask.array as da\n",
    "\n",
    "x = da.random.random((20_000, 20_000), chunks=(5_000, 1_000))\n",
    "result = (x + x.T).mean(axis=0).sum()\n",
    "result.compute()\n",
    "```"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "46fd60d2-87c0-47df-a857-a37cf6ee26f7",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Your solution to Exercise 2 here"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "628ceef7-8602-4e13-8191-02d6078cdcaf",
   "metadata": {
    "jupyter": {
     "source_hidden": true
    },
    "scrolled": true,
    "tags": []
   },
   "outputs": [],
   "source": [
    "# Solution to Exercise 2\n",
    "import time\n",
    "\n",
    "class TaskTimerPlugin(SchedulerPlugin):\n",
    "    def __init__(self):\n",
    "        self.start_times = {}\n",
    "        self.stop_times = {}\n",
    "        self.task_durations = {}\n",
    "\n",
    "    def transition(self, key, start, finish, *args, **kwargs):\n",
    "        if finish == \"processing\": \n",
    "            self.start_times[key] = time.time()\n",
    "        elif start ==\"processing\" and finish == \"memory\":\n",
    "            self.stop_times[key] = time.time()\n",
    "            self.task_durations[key] = self.stop_times[key] - self.start_times[key]\n",
    "\n",
    "# Create LocalCluster and Client\n",
    "cluster = LocalCluster()\n",
    "client = Client(cluster)\n",
    "\n",
    "# Instantiate and add the TaskTimerPlugin to our cluster\n",
    "plugin = TaskTimerPlugin()\n",
    "cluster.scheduler.add_plugin(plugin)\n",
    "\n",
    "import dask.array as da\n",
    "\n",
    "x = da.random.random((20_000, 20_000), chunks=(5_000, 1_000))\n",
    "result = (x + x.T).mean(axis=0).sum()\n",
    "result.compute()\n",
    "\n",
    "plugin.task_durations"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9e230a95-353e-4d7a-a5b1-16f8b6e97b9f",
   "metadata": {},
   "source": [
    "**Bonus**: If you have extra time, make a plot of the task duration distribution (hint: `pandas` and `matplotlib` are installed)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3ab78fb2-ef32-4587-a0fe-0b4c31930926",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "# Your plotting code here"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d5369156-5524-49c0-bb00-37adffa1d175",
   "metadata": {
    "jupyter": {
     "source_hidden": true
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "#solution\n",
    "import pandas as pd\n",
    "\n",
    "df = pd.DataFrame([(key, 1_000 * value) for key, value in plugin.task_durations.items()],\n",
    "                  columns=[\"key\", \"duration\"])\n",
    "ax = df.duration.plot(kind=\"hist\", bins=50, logy=True)\n",
    "ax.set_xlabel(\"Task duration [ms]\")\n",
    "ax.set_ylabel(\"Counts\");"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5bf57a9e-8fb6-4355-9aff-3f426efab9ce",
   "metadata": {},
   "outputs": [],
   "source": [
    "client.close()\n",
    "cluster.close()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f09c76dd-7523-42f5-a483-30d8ba69f2ec",
   "metadata": {},
   "source": [
    "# Summary\n",
    "\n",
    "This notebook we took a detailed look at the components of a Dask cluster, illustrated how to inspect the internal state of a cluster (both the scheduler and workers), and how you can use Dask's plugin system to execute custom code during a cluster's lifecycle.\n"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.11.4"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}