3-graph-optimization.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "recovered-trouble",
   "metadata": {},
   "source": [
    "# Advanced Collections: _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\n",
    "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",
    "  - [Graph Optimization](https://docs.dask.org/en/latest/optimization.html)\n",
    "  \n",
    "The ``dask.optimization`` module contains several functions to transform graphs\n",
    "in a variety of useful ways. In most cases, users won't need to interact with\n",
    "these functions directly, as specialized subsets of these transforms are done\n",
    "automatically in the Dask collections (``dask.array``, ``dask.bag``, and\n",
    "``dask.dataframe``). However, users working with custom graphs or computations\n",
    "may find that applying these methods results in substantial speedups.\n",
    "\n",
    "## Example\n",
    "\n",
    "Suppose you had a custom Dask graph for doing a word counting task.\n",
    "\n",
    "> **NOTE**: To gain intuition about optimization we'll be looking at just the task graph first. We will talk about how this relates to collections at the end.\n",
    "\n",
    "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"
   ]
  },
  {
   "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)"
   ]
  },
  {
   "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)"
   ]
  },
  {
   "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."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "utility-document",
   "metadata": {},
   "source": [
    "#### In practice\n",
    "\n",
    "Cull is very useful for operating on partitioned data. For instance consider the\n",
    "case of a timeseries stored in a parquet file where each parition represents one\n",
    "day. Let's generate such a timeseries:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "thermal-symphony",
   "metadata": {},
   "outputs": [],
   "source": [
    "import dask\n",
    "import dask.dataframe as dd\n",
    "\n",
    "dask.datasets.timeseries().to_parquet(\"timeseries\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "sophisticated-asset",
   "metadata": {},
   "source": [
    "Now read in the data and look at the task graph:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "amber-locator",
   "metadata": {},
   "outputs": [],
   "source": [
    "ddf = dd.read_parquet(\"timeseries\")\n",
    "ddf.visualize(optimize_graph=True)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "laden-trout",
   "metadata": {},
   "source": [
    "If you select a subset of the days, and optimize the graph you can see that there is now only one task at the bottom of the task graph. That is because Dask knows that that all the other partitions don't have data that you need, so those tasks have been culled."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "aerial-choice",
   "metadata": {},
   "outputs": [],
   "source": [
    "ddf[\"2000-01-13\"].visualize(optimize_graph=True)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "integral-handling",
   "metadata": {},
   "source": [
    "### Inline\n",
    "\n",
    "Back to the fruits!\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)"
   ]
  },
  {
   "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)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "numerical-contrast",
   "metadata": {},
   "source": [
    "### Fuse\n",
    "\n",
    "Now we have a set of purely linear tasks. We'd like to have the scheduler run\n",
    "all of these on the same worker to reduce data serialization between workers.\n",
    "One option is just to merge these linear chains into one big task using the\n",
    "``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)"
   ]
  },
  {
   "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)"
   ]
  },
  {
   "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."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "stone-surfing",
   "metadata": {},
   "source": [
    "## Customizing Optimization\n",
    "\n",
    "Dask defines a default optimization strategy for each collection type (Array,\n",
    "Bag, DataFrame, Delayed).  However, different applications may have different\n",
    "needs.  To address this variability of needs, you can construct your own custom\n",
    "optimization function and use it instead of the default.  An optimization\n",
    "function takes in a task graph and list of desired keys and returns a new\n",
    "task graph:\n",
    "\n",
    "```python\n",
    "def my_optimize_function(dsk, keys):\n",
    "    new_dsk = {...}\n",
    "    return new_dsk\n",
    "```\n",
    "\n",
    "You can then register this optimization class against whichever collection type\n",
    "you prefer and it will be used instead of the default scheme:\n",
    "\n",
    "```python\n",
    "with dask.config.set(array_optimize=my_optimize_function):\n",
    "    x, y = dask.compute(x, y)\n",
    "```\n",
    "\n",
    "You can register separate optimization functions for different collections, or\n",
    "you can register ``None`` if you do not want particular types of collections to\n",
    "be optimized:\n",
    "\n",
    "```python\n",
    "with dask.config.set(array_optimize=my_optimize_function,\n",
    "                     dataframe_optimize=None,\n",
    "                     delayed_optimize=my_other_optimize_function):\n",
    "    ...\n",
    "```\n",
    "\n",
    "You do not need to specify all collections.  Collections will default to their\n",
    "standard optimization scheme (which is usually a good choice).\n",
    "\n",
    "When creating your own collection, you can specify the optimization function\n",
    "by implementing a ``_dask_optimize__`` method."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "identified-brain",
   "metadata": {},
   "source": [
    "## Exercise\n",
    "\n",
    "This exercise is based off of [graphchain](https://github.com/radix-ai/graphchain)\n",
    "which implements a custom optimization function that reads data from a\n",
    "cache if it has already been computed.\n",
    "\n",
    "Take a look at the following task graph."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "massive-pasta",
   "metadata": {},
   "outputs": [],
   "source": [
    "import dask\n",
    "import pandas as pd\n",
    "\n",
    "def create_dataframe(num_rows, num_cols):\n",
    "    print('Creating DataFrame...')\n",
    "    return pd.DataFrame(data=[range(num_cols)]*num_rows)\n",
    "\n",
    "def complicated_computation(df, num_quantiles):\n",
    "    print('Running complicated computation on DataFrame...')\n",
    "    return df.quantile(q=[i / num_quantiles for i in range(num_quantiles)])\n",
    "\n",
    "def summarise_dataframes(*dfs):\n",
    "    print('Summing DataFrames...')\n",
    "    return sum(df.sum().sum() for df in dfs)\n",
    "\n",
    "dsk = {\n",
    "    'df_a': (create_dataframe, 10_000, 1000),\n",
    "    'df_b': (create_dataframe, 10_000, 1000),\n",
    "    'df_c': (complicated_computation, 'df_a', 2048),\n",
    "    'df_d': (complicated_computation, 'df_b', 2048),\n",
    "    'result': (summarise_dataframes, 'df_c', 'df_d')\n",
    "}\n",
    "\n",
    "%time dask.get(dsk, 'result')"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "earned-sight",
   "metadata": {},
   "source": [
    "Notice how `df_a` and `df_b` have the same definition?\n",
    "\n",
    "Write an optimization function that replaces one of the matching tasks with a key\n",
    "pointing to the other task.\n",
    "\n",
    "> NOTE: This is a much more simplistic approach than graphchain uses, and it's less\n",
    "> useful in a distributed context."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "expired-mandate",
   "metadata": {},
   "outputs": [],
   "source": [
    "def optimize(dsk):\n",
    "    new_dsk = {}  # make sure that you don't alter the task graph inplace\n",
    "    \n",
    "    for k, v in dsk.items():\n",
    "        ... # if any two tasks match, replace one of the matching tasks with a key pointing to the other task.\n",
    "    return new_dsk\n",
    "\n",
    "\n",
    "%time dask.get(optimize(dsk), 'result')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "bound-specific",
   "metadata": {
    "jupyter": {
     "source_hidden": true
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "#solution\n",
    "def optimize(dsk):\n",
    "    new_dsk = {}\n",
    "    \n",
    "    for k, v in dsk.items():\n",
    "        if v not in new_dsk.values():\n",
    "            new_dsk[k] = v\n",
    "        else:\n",
    "            new_dsk[k] = {v: k for k, v in new_dsk.items()}[v]\n",
    "    return new_dsk\n",
    "\n",
    "%time dask.get(optimize(dsk), 'result')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "relevant-punch",
   "metadata": {
    "jupyter": {
     "source_hidden": true
    }
   },
   "outputs": [],
   "source": [
    "# alternate solution\n",
    "def optimize(dsk):\n",
    "    new_dsk = {}\n",
    "\n",
    "    for k, v in dsk.items():\n",
    "        first_key = list(dsk.keys())[list(dsk.values()).index(v)]\n",
    "        if first_key not in new_dsk.keys():\n",
    "            new_dsk[k] = v\n",
    "        else:\n",
    "            new_dsk[k] = first_key\n",
    "    return new_dsk\n",
    "\n",
    "%time dask.get(optimize(dsk), 'result')"
   ]
  },
  {
   "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",
   "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.9.4"
  },
  "toc-autonumbering": false,
  "toc-showcode": false,
  "toc-showmarkdowntxt": false,
  "toc-showtags": false
 },
 "nbformat": 4,
 "nbformat_minor": 5
}