appendix-graph-optimization.ipynb

{
 "cells": [
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "recovered-trouble",
   "metadata": {},
   "source": [
    "<img src=\"images/dask_horizontal.svg\" width=\"45%\" alt=\"Dask logo\\\">\n",
    "\n",
    "# Graph Optimizations\n",
    "\n",
    "In general, there are two goals when doing graph optimizations:\n",
    "\n",
    "1. Simplify computation\n",
    "2. Improve parallelism\n",
    "\n",
    "Simplifying computation can be done on a graph level by removing unnecessary tasks (``cull``).\n",
    "\n",
    "Parallelism can be improved by reducing\n",
    "inter-task communication, whether by fusing many tasks into one (``fuse``), or\n",
    "by inlining cheap operations (``inline``, ``inline_functions``).\n",
    "\n",
    "\n",
    "**Related Documentation**\n",
    "\n",
    "  - [Optimization](https://docs.dask.org/en/latest/optimize.html)\n",
    "\n",
    "## Example\n",
    "\n",
    "Suppose you had a custom Dask graph for doing a word counting task:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "disabled-greek",
   "metadata": {},
   "outputs": [],
   "source": [
    "import dask\n",
    "\n",
    "def print_and_return(string):\n",
    "    print(string)\n",
    "    return string\n",
    "\n",
    "def format_str(count, val, nwords):\n",
    "    return (f'word list has {count} occurrences of '\n",
    "            f'{val}, out of {nwords} words')\n",
    "\n",
    "dsk = {'words': 'apple orange apple pear orange pear pear',\n",
    "       'nwords': (len, (str.split, 'words')),\n",
    "       'val1': 'orange',\n",
    "       'val2': 'apple',\n",
    "       'val3': 'pear',\n",
    "       'count1': (str.count, 'words', 'val1'),\n",
    "       'count2': (str.count, 'words', 'val2'),\n",
    "       'count3': (str.count, 'words', 'val3'),\n",
    "       'format1': (format_str, 'count1', 'val1', 'nwords'),\n",
    "       'format2': (format_str, 'count2', 'val2', 'nwords'),\n",
    "       'format3': (format_str, 'count3', 'val3', 'nwords'),\n",
    "       'print1': (print_and_return, 'format1'),\n",
    "       'print2': (print_and_return, 'format2'),\n",
    "       'print3': (print_and_return, 'format3'),\n",
    "}\n",
    "\n",
    "dask.visualize(dsk, verbose=True, collapse_outputs=True)"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "da3a382a",
   "metadata": {},
   "source": [
    "In this example we are:\n",
    "\n",
    "1. counting the frequency of the words ``'orange'``, ``'apple'``, and ``'pear'`` in the list of words\n",
    "2. formatting an output string reporting the results\n",
    "3. printing the output and returning the output string"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "circular-driving",
   "metadata": {},
   "source": [
    "### Cull\n",
    "\n",
    "To perform the computation, we first remove unnecessary components from the\n",
    "graph using the ``cull`` function and then pass the Dask graph and the desired\n",
    "output keys to a scheduler ``get`` function:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "meaningful-knife",
   "metadata": {},
   "outputs": [],
   "source": [
    "from dask.threaded import get\n",
    "from dask.optimization import cull\n",
    "\n",
    "outputs = ['print1', 'print2']\n",
    "dsk1, dependencies = cull(dsk, outputs)  # remove unnecessary tasks from the graph\n",
    "\n",
    "results = get(dsk1, outputs)\n",
    "dask.visualize(dsk1, verbose=True, collapse_outputs=True)"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "afraid-asbestos",
   "metadata": {},
   "source": [
    "As can be seen above, the scheduler computed only the requested outputs\n",
    "(``'print3'`` was never computed). This is because we called the\n",
    "``dask.optimization.cull`` function, which removes the unnecessary tasks from\n",
    "the graph.\n",
    "\n",
    "Culling is part of the default optimization pass of almost all collections.\n",
    "Often you want to call it somewhat early to reduce the amount of work done in\n",
    "later steps."
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "integral-handling",
   "metadata": {},
   "source": [
    "### Inline\n",
    "\n",
    "Looking at the word counting task graph, there are multiple accesses to constants such\n",
    "as ``'val1'`` or ``'val2'``. These can be inlined into the\n",
    "tasks to improve efficiency using the ``inline`` function. For example:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "lightweight-gamma",
   "metadata": {},
   "outputs": [],
   "source": [
    "from dask.optimization import inline\n",
    "\n",
    "dsk2 = inline(dsk1, dependencies=dependencies)\n",
    "results = get(dsk2, outputs)\n",
    "dask.visualize(dsk2, verbose=True, collapse_outputs=True)"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "promising-retro",
   "metadata": {},
   "source": [
    "Now we have two sets of *almost* linear task chains. The only link between them\n",
    "is the word counting function. For cheap operations like this, the\n",
    "serialization cost may be larger than the actual computation, so it may be\n",
    "faster to do the computation more than once, rather than passing the results to\n",
    "all nodes. To perform this function inlining, the ``inline_functions`` function\n",
    "can be used:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "stretch-doctor",
   "metadata": {},
   "outputs": [],
   "source": [
    "from dask.optimization import inline_functions\n",
    "\n",
    "dsk3 = inline_functions(dsk2, outputs, [len, str.split],\n",
    "                        dependencies=dependencies)\n",
    "results = get(dsk3, outputs)\n",
    "dask.visualize(dsk3, verbose=True, collapse_outputs=True)"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "5e54f219",
   "metadata": {},
   "source": [
    "Now we have a set of purely linear tasks. We’d like to have the scheduler run all of these on the same worker to reduce data serialization between workers. One option is just to merge these linear chains into one big task using the ``fuse`` function:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "cleared-shoulder",
   "metadata": {},
   "outputs": [],
   "source": [
    "from dask.optimization import fuse\n",
    "\n",
    "dsk4, dependencies = fuse(dsk3)\n",
    "results = get(dsk4, outputs)\n",
    "dask.visualize(dsk4, verbose=True, collapse_outputs=True)"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "weird-blade",
   "metadata": {},
   "source": [
    "### Result\n",
    "\n",
    "Putting it all together:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "experimental-special",
   "metadata": {},
   "outputs": [],
   "source": [
    "def optimize_and_get(dsk, keys):\n",
    "    dsk1, deps = cull(dsk, keys)\n",
    "    dsk2 = inline(dsk1, dependencies=deps)\n",
    "    dsk3 = inline_functions(dsk2, keys, [len, str.split],\n",
    "                            dependencies=deps)\n",
    "    dsk4, deps = fuse(dsk3)\n",
    "    return get(dsk4, keys)\n",
    "\n",
    "optimize_and_get(dsk, outputs)"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "hearing-rochester",
   "metadata": {},
   "source": [
    "In summary, the above operations accomplish the following:\n",
    "\n",
    "1. Removed tasks unnecessary for the desired output using ``cull``\n",
    "2. Inlined constants using ``inline``\n",
    "3. Inlined cheap computations using ``inline_functions``, improving parallelism\n",
    "4. Fused linear tasks together to ensure they run on the same worker using ``fuse``\n",
    "\n",
    "These optimizations are already performed automatically in the Dask collections."
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "b222c529-d522-47ec-853d-d23d2744b7b1",
   "metadata": {},
   "source": [
    "## Conclusion\n",
    "\n",
    "Optimizations in Dask let's you simplify computation and improve parallelism. There are some great ones included by default (`cull`, `inline`, `fuse`), but sometimes it can be really powerful to write custom optimizations and either use them on existing collections or on custom collections. We'll touch on this a bit more in the next section about custom collections."
   ]
  }
 ],
 "metadata": {
  "jupytext": {
   "cell_metadata_filter": "-all",
   "main_language": "python",
   "notebook_metadata_filter": "-all",
   "text_representation": {
    "extension": ".md",
    "format_name": "markdown"
   }
  },
  "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.11"
  },
  "toc-autonumbering": false,
  "toc-showcode": false,
  "toc-showmarkdowntxt": false,
  "toc-showtags": false
 },
 "nbformat": 4,
 "nbformat_minor": 5
}