dask_basics.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "08476042-7d6d-43e9-bf61-20f49090fcde",
   "metadata": {},
   "source": [
    "# Dask Basics Example\n",
    "\n",
    "This notebook goes through some basic dask commands to run some simple (non-coffea) computations on the cluster."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "f1ce613c-4f5f-48ef-9bc1-9bcfc410fdab",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "from dask.distributed import Client, progress"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c8e983ea-1907-4c59-90ae-c51f9f8565d5",
   "metadata": {},
   "source": [
    "This next cell sets up a Condor cluster object, which is where we can make any ClassAd specifications we need to make. We'll also point an environment variable to some necessary Condor configurations.\n",
    "\n",
    "Near the end of the cell, we ask for the cluster to give us between 1 and 4 threads at a time when we perform computations."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "6b722d96-5452-4954-a285-95e2c003784e",
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "data": {
      "text/html": [
       "<div>\n",
       "    <div style=\"width: 24px; height: 24px; background-color: #e1e1e1; border: 3px solid #9D9D9D; border-radius: 5px; position: absolute;\"> </div>\n",
       "    <div style=\"margin-left: 48px;\">\n",
       "        <h3 style=\"margin-bottom: 0px;\">Client</h3>\n",
       "        <p style=\"color: #9D9D9D; margin-bottom: 0px;\">Client-4863e211-bef7-11ee-a456-3cecefda85c8</p>\n",
       "        <table style=\"width: 100%; text-align: left;\">\n",
       "\n",
       "        <tr>\n",
       "        \n",
       "            <td style=\"text-align: left;\"><strong>Connection method:</strong> Cluster object</td>\n",
       "            <td style=\"text-align: left;\"><strong>Cluster type:</strong> dask_jobqueue.HTCondorCluster</td>\n",
       "        \n",
       "        </tr>\n",
       "\n",
       "        \n",
       "            <tr>\n",
       "                <td style=\"text-align: left;\">\n",
       "                    <strong>Dashboard: </strong> <a href=\"proxy/8787/status\" target=\"_blank\">proxy/8787/status</a>\n",
       "                </td>\n",
       "                <td style=\"text-align: left;\"></td>\n",
       "            </tr>\n",
       "        \n",
       "\n",
       "        </table>\n",
       "\n",
       "        \n",
       "            <button style=\"margin-bottom: 12px;\" data-commandlinker-command=\"dask:populate-and-launch-layout\" data-commandlinker-args='{\"url\": \"proxy/8787/status\" }'>\n",
       "                Launch dashboard in JupyterLab\n",
       "            </button>\n",
       "        \n",
       "\n",
       "        \n",
       "            <details>\n",
       "            <summary style=\"margin-bottom: 20px;\"><h3 style=\"display: inline;\">Cluster Info</h3></summary>\n",
       "            <div class=\"jp-RenderedHTMLCommon jp-RenderedHTML jp-mod-trusted jp-OutputArea-output\">\n",
       "    <div style=\"width: 24px; height: 24px; background-color: #e1e1e1; border: 3px solid #9D9D9D; border-radius: 5px; position: absolute;\">\n",
       "    </div>\n",
       "    <div style=\"margin-left: 48px;\">\n",
       "        <h3 style=\"margin-bottom: 0px; margin-top: 0px;\">HTCondorCluster</h3>\n",
       "        <p style=\"color: #9D9D9D; margin-bottom: 0px;\">b6353852</p>\n",
       "        <table style=\"width: 100%; text-align: left;\">\n",
       "            <tr>\n",
       "                <td style=\"text-align: left;\">\n",
       "                    <strong>Dashboard:</strong> <a href=\"proxy/8787/status\" target=\"_blank\">proxy/8787/status</a>\n",
       "                </td>\n",
       "                <td style=\"text-align: left;\">\n",
       "                    <strong>Workers:</strong> 0\n",
       "                </td>\n",
       "            </tr>\n",
       "            <tr>\n",
       "                <td style=\"text-align: left;\">\n",
       "                    <strong>Total threads:</strong> 0\n",
       "                </td>\n",
       "                <td style=\"text-align: left;\">\n",
       "                    <strong>Total memory:</strong> 0 B\n",
       "                </td>\n",
       "            </tr>\n",
       "            \n",
       "        </table>\n",
       "\n",
       "        <details>\n",
       "            <summary style=\"margin-bottom: 20px;\">\n",
       "                <h3 style=\"display: inline;\">Scheduler Info</h3>\n",
       "            </summary>\n",
       "\n",
       "            <div style=\"\">\n",
       "    <div>\n",
       "        <div style=\"width: 24px; height: 24px; background-color: #FFF7E5; border: 3px solid #FF6132; border-radius: 5px; position: absolute;\"> </div>\n",
       "        <div style=\"margin-left: 48px;\">\n",
       "            <h3 style=\"margin-bottom: 0px;\">Scheduler</h3>\n",
       "            <p style=\"color: #9D9D9D; margin-bottom: 0px;\">Scheduler-b9b008cf-6d31-4bc7-8fe3-b17fd1db7b16</p>\n",
       "            <table style=\"width: 100%; text-align: left;\">\n",
       "                <tr>\n",
       "                    <td style=\"text-align: left;\">\n",
       "                        <strong>Comm:</strong> tcp://144.92.181.248:27997\n",
       "                    </td>\n",
       "                    <td style=\"text-align: left;\">\n",
       "                        <strong>Workers:</strong> 0\n",
       "                    </td>\n",
       "                </tr>\n",
       "                <tr>\n",
       "                    <td style=\"text-align: left;\">\n",
       "                        <strong>Dashboard:</strong> <a href=\"proxy/8787/status\" target=\"_blank\">proxy/8787/status</a>\n",
       "                    </td>\n",
       "                    <td style=\"text-align: left;\">\n",
       "                        <strong>Total threads:</strong> 0\n",
       "                    </td>\n",
       "                </tr>\n",
       "                <tr>\n",
       "                    <td style=\"text-align: left;\">\n",
       "                        <strong>Started:</strong> Just now\n",
       "                    </td>\n",
       "                    <td style=\"text-align: left;\">\n",
       "                        <strong>Total memory:</strong> 0 B\n",
       "                    </td>\n",
       "                </tr>\n",
       "            </table>\n",
       "        </div>\n",
       "    </div>\n",
       "\n",
       "    <details style=\"margin-left: 48px;\">\n",
       "        <summary style=\"margin-bottom: 20px;\">\n",
       "            <h3 style=\"display: inline;\">Workers</h3>\n",
       "        </summary>\n",
       "\n",
       "        \n",
       "\n",
       "    </details>\n",
       "</div>\n",
       "\n",
       "        </details>\n",
       "    </div>\n",
       "</div>\n",
       "            </details>\n",
       "        \n",
       "\n",
       "    </div>\n",
       "</div>"
      ],
      "text/plain": [
       "<Client: 'tcp://144.92.181.248:27997' processes=0 threads=0, memory=0 B>"
      ]
     },
     "execution_count": 2,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# this cell starts workers in condor\n",
    "import os\n",
    "from dask_jobqueue import HTCondorCluster\n",
    "from dask import delayed\n",
    "from dask.distributed import Client, as_completed\n",
    "\n",
    "os.environ[\"CONDOR_CONFIG\"] = \"/etc/condor/condor_config\"\n",
    "\n",
    "cluster = HTCondorCluster(\n",
    "    cores=1,\n",
    "    memory='4gb',\n",
    "    disk='1 GB',\n",
    "    job_extra_directives={\n",
    "       \"+SingularityImage\" : '\"docker://index.docker.io/coffeateam/coffea-dask:0.7.22-py3.10-g7cbcc\"',\n",
    "        \"Requirements\": \"HasSingularityJobStart\",\n",
    "    }\n",
    ")\n",
    "cluster.adapt(minimum=1, maximum=4)\n",
    "\n",
    "\n",
    "client = Client(cluster)\n",
    "client"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ed187cb6-d567-4e0d-8f92-da3b145d12fc",
   "metadata": {},
   "source": [
    "The above cell should have created a client object, with a button that says \"Launch dashboard in JupyterLab\". Clicking that button allows you to watch the Condor workers' memory usage and activity.\n",
    "\n",
    "Now, we'll define some functions and run them locally."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "d5523bf3-a35b-4b68-be1c-31610b87b4e6",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "import time\n",
    "import random\n",
    "\n",
    "def inc(x):\n",
    "    time.sleep(random.random())\n",
    "    return x + 1\n",
    "\n",
    "def dec(x):\n",
    "    time.sleep(random.random())\n",
    "    return x - 1\n",
    "\n",
    "def add(x, y):\n",
    "    time.sleep(random.random())\n",
    "    return x + y"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "84bf82b8-cc93-4a4d-85a5-bc5116bdf381",
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "CPU times: user 14.7 ms, sys: 92.1 ms, total: 107 ms\n",
      "Wall time: 741 ms\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "3"
      ]
     },
     "execution_count": 4,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "%%time\n",
    "x = inc(1)\n",
    "y = dec(2)\n",
    "z = add(x, y)\n",
    "z"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7471e9e6-7258-46d1-93b2-7feee1e4a0a5",
   "metadata": {},
   "source": [
    "Now, we'll \"Dask-ify\" these functions using `dask.delayed`. Calling these functions does not lead to immediate (\"eager\") computations, but instead waits until we call `compute` on the functions' result to perform the calculation. This is automatically done on the cluster we created."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "b608dffb-9365-401d-9e30-db27ead2583f",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "import dask\n",
    "inc = dask.delayed(inc)\n",
    "dec = dask.delayed(dec)\n",
    "add = dask.delayed(add)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "1ac72649-341c-4444-9103-48effc131d6c",
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "Delayed('add-5de68405-a152-4cbb-ad51-bf9895dfcbe3')"
      ]
     },
     "execution_count": 6,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "x = inc(1)\n",
    "y = dec(2)\n",
    "z = add(x, y)\n",
    "z"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "20bca461-96fb-4fb3-a703-20b213d29b97",
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "3"
      ]
     },
     "execution_count": 7,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "z.compute()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a5e5b455-eac9-4ad6-8707-5691751297b9",
   "metadata": {},
   "source": [
    "Now, we'll create a whole list of computations to perform (256 of them). Instead of using `compute`, we'll first call `persist`, since we're working with a large array. This triggers computation in the background, but keeps the computation on the workers until it is actually called with `compute`, typically once reduced to a smaller value."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "3da3d348-174a-4ef2-9591-62ea0eee4be2",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "zs = []"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "b7eacfb0-6ce4-4d84-8a6a-584ca3476746",
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "CPU times: user 39.8 ms, sys: 16.5 ms, total: 56.4 ms\n",
      "Wall time: 46.5 ms\n"
     ]
    }
   ],
   "source": [
    "%%time\n",
    "for i in range(256):\n",
    "    x = inc(i)\n",
    "    y = dec(x)\n",
    "    z = add(x, y)\n",
    "    zs.append(z)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "5919e6f4-3d37-4c85-9555-241312930891",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "zs = dask.persist(*zs)  # trigger computation in the background"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "896a9652-2167-4180-bd83-c9f4566433ba",
   "metadata": {},
   "source": [
    "`zs` is now a list of \"futures\" objects on the cluster. We can then perform new Dask \"lazy\" computations on those (in the while loop here) and then call `compute` once we've reduced the calculation to a single number.\n",
    "\n",
    "First, we'll ask for ten workers to try and get more jobs running at a time."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "id": "e43739c7-226d-497c-af96-e4f4766a2886",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "client.cluster.scale(10)  # ask for ten 4-thread workers"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "id": "a65c5a48-9771-4d5c-98d8-5a5e539f2c96",
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "([65536],)"
      ]
     },
     "execution_count": 15,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "L = zs\n",
    "while len(L) > 1:\n",
    "    new_L = []\n",
    "    for i in range(0, len(L), 2):\n",
    "        lazy = add(L[i], L[i + 1])  # add neighbors\n",
    "        new_L.append(lazy)\n",
    "    L = new_L                       # swap old list for new\n",
    "\n",
    "dask.compute(L)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6eb8db6c-b9d5-4d12-b4df-d2a188c4559a",
   "metadata": {},
   "source": [
    "At the end, it's good to shut down the cluster to avoid issues."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "id": "cea419d0-8a6f-424c-a3f6-40f221cb4a4f",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "# shutdown distributed.localCluster\n",
    "client.shutdown()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ef81d0a0-ecf9-4974-9fac-6f4176b140a2",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "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.10.8"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}