dask.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "e55ef8ce-8960-439c-9b34-6d86dad8a3d0",
   "metadata": {
    "editable": true,
    "slideshow": {
     "slide_type": "slide"
    },
    "tags": []
   },
   "source": [
    "# Dask: a flexible library for parallel computing in Python."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f8744560-cdd4-4e86-8d1b-c65be1abf177",
   "metadata": {
    "editable": true,
    "slideshow": {
     "slide_type": "fragment"
    },
    "tags": []
   },
   "source": [
    "Dask is composed of two parts:\n",
    "\n",
    "- Dynamic task scheduling optimized for computation. This is similar to Airflow, Luigi, Celery, or Make, but optimized for interactive computational workloads.\n",
    "- “Big Data” collections like parallel arrays, dataframes, and lists that extend common interfaces like NumPy, Pandas, or Python iterators to larger-than-memory or distributed environments. These parallel collections run on top of dynamic task schedulers."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a1ad8f97-aa74-495d-9055-423e12935634",
   "metadata": {
    "editable": true,
    "slideshow": {
     "slide_type": "fragment"
    },
    "tags": []
   },
   "source": [
    "Many classes of problems cannot fit on a single computer or compute node. \n",
    "Some require more memory to process very large datasets, or could simply benefit from being spread across multiple compute nodes. \n",
    "One tool that can span from laptops to large clusters is `dask` (https://docs.dask.org/en/stable/), which can be installed via `conda`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 28,
   "id": "a38bf62e-2df7-4911-84b7-1ec9d7a84f72",
   "metadata": {
    "editable": true,
    "slideshow": {
     "slide_type": "subslide"
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "import dask.dataframe as dd\n",
    "import dask.array as da"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 29,
   "id": "4fa5b99b-ddb3-4467-b036-6de955103e6f",
   "metadata": {
    "editable": true,
    "slideshow": {
     "slide_type": "fragment"
    },
    "tags": []
   },
   "outputs": [
    {
     "data": {
      "text/html": [
       "<table>\n",
       "    <tr>\n",
       "        <td>\n",
       "            <table style=\"border-collapse: collapse;\">\n",
       "                <thead>\n",
       "                    <tr>\n",
       "                        <td> </td>\n",
       "                        <th> Array </th>\n",
       "                        <th> Chunk </th>\n",
       "                    </tr>\n",
       "                </thead>\n",
       "                <tbody>\n",
       "                    \n",
       "                    <tr>\n",
       "                        <th> Bytes </th>\n",
       "                        <td> 781.25 kiB </td>\n",
       "                        <td> 78.12 kiB </td>\n",
       "                    </tr>\n",
       "                    \n",
       "                    <tr>\n",
       "                        <th> Shape </th>\n",
       "                        <td> (200, 500) </td>\n",
       "                        <td> (100, 100) </td>\n",
       "                    </tr>\n",
       "                    <tr>\n",
       "                        <th> Dask graph </th>\n",
       "                        <td colspan=\"2\"> 10 chunks in 1 graph layer </td>\n",
       "                    </tr>\n",
       "                    <tr>\n",
       "                        <th> Data type </th>\n",
       "                        <td colspan=\"2\"> float64 numpy.ndarray </td>\n",
       "                    </tr>\n",
       "                </tbody>\n",
       "            </table>\n",
       "        </td>\n",
       "        <td>\n",
       "        <svg width=\"170\" height=\"98\" style=\"stroke:rgb(0,0,0);stroke-width:1\" >\n",
       "\n",
       "  <!-- Horizontal lines -->\n",
       "  <line x1=\"0\" y1=\"0\" x2=\"120\" y2=\"0\" style=\"stroke-width:2\" />\n",
       "  <line x1=\"0\" y1=\"24\" x2=\"120\" y2=\"24\" />\n",
       "  <line x1=\"0\" y1=\"48\" x2=\"120\" y2=\"48\" style=\"stroke-width:2\" />\n",
       "\n",
       "  <!-- Vertical lines -->\n",
       "  <line x1=\"0\" y1=\"0\" x2=\"0\" y2=\"48\" style=\"stroke-width:2\" />\n",
       "  <line x1=\"24\" y1=\"0\" x2=\"24\" y2=\"48\" />\n",
       "  <line x1=\"48\" y1=\"0\" x2=\"48\" y2=\"48\" />\n",
       "  <line x1=\"72\" y1=\"0\" x2=\"72\" y2=\"48\" />\n",
       "  <line x1=\"96\" y1=\"0\" x2=\"96\" y2=\"48\" />\n",
       "  <line x1=\"120\" y1=\"0\" x2=\"120\" y2=\"48\" style=\"stroke-width:2\" />\n",
       "\n",
       "  <!-- Colored Rectangle -->\n",
       "  <polygon points=\"0.0,0.0 120.0,0.0 120.0,48.0 0.0,48.0\" style=\"fill:#ECB172A0;stroke-width:0\"/>\n",
       "\n",
       "  <!-- Text -->\n",
       "  <text x=\"60.000000\" y=\"68.000000\" font-size=\"1.0rem\" font-weight=\"100\" text-anchor=\"middle\" >500</text>\n",
       "  <text x=\"140.000000\" y=\"24.000000\" font-size=\"1.0rem\" font-weight=\"100\" text-anchor=\"middle\" transform=\"rotate(-90,140.000000,24.000000)\">200</text>\n",
       "</svg>\n",
       "        </td>\n",
       "    </tr>\n",
       "</table>"
      ],
      "text/plain": [
       "dask.array<array, shape=(200, 500), dtype=float64, chunksize=(100, 100), chunktype=numpy.ndarray>"
      ]
     },
     "execution_count": 29,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "data = np.random.normal(size=100000).reshape(200, 500)\n",
    "a = da.from_array(data, chunks=(100, 100))\n",
    "a"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 30,
   "id": "3808a3f2-0c36-4b1e-b91d-4937935f21c7",
   "metadata": {
    "editable": true,
    "slideshow": {
     "slide_type": "subslide"
    },
    "tags": []
   },
   "outputs": [
    {
     "data": {
      "text/html": [
       "<table>\n",
       "    <tr>\n",
       "        <td>\n",
       "            <table style=\"border-collapse: collapse;\">\n",
       "                <thead>\n",
       "                    <tr>\n",
       "                        <td> </td>\n",
       "                        <th> Array </th>\n",
       "                        <th> Chunk </th>\n",
       "                    </tr>\n",
       "                </thead>\n",
       "                <tbody>\n",
       "                    \n",
       "                    <tr>\n",
       "                        <th> Bytes </th>\n",
       "                        <td> 400 B </td>\n",
       "                        <td> 400 B </td>\n",
       "                    </tr>\n",
       "                    \n",
       "                    <tr>\n",
       "                        <th> Shape </th>\n",
       "                        <td> (50,) </td>\n",
       "                        <td> (50,) </td>\n",
       "                    </tr>\n",
       "                    <tr>\n",
       "                        <th> Dask graph </th>\n",
       "                        <td colspan=\"2\"> 1 chunks in 2 graph layers </td>\n",
       "                    </tr>\n",
       "                    <tr>\n",
       "                        <th> Data type </th>\n",
       "                        <td colspan=\"2\"> float64 numpy.ndarray </td>\n",
       "                    </tr>\n",
       "                </tbody>\n",
       "            </table>\n",
       "        </td>\n",
       "        <td>\n",
       "        <svg width=\"170\" height=\"79\" style=\"stroke:rgb(0,0,0);stroke-width:1\" >\n",
       "\n",
       "  <!-- Horizontal lines -->\n",
       "  <line x1=\"0\" y1=\"0\" x2=\"120\" y2=\"0\" style=\"stroke-width:2\" />\n",
       "  <line x1=\"0\" y1=\"29\" x2=\"120\" y2=\"29\" style=\"stroke-width:2\" />\n",
       "\n",
       "  <!-- Vertical lines -->\n",
       "  <line x1=\"0\" y1=\"0\" x2=\"0\" y2=\"29\" style=\"stroke-width:2\" />\n",
       "  <line x1=\"120\" y1=\"0\" x2=\"120\" y2=\"29\" style=\"stroke-width:2\" />\n",
       "\n",
       "  <!-- Colored Rectangle -->\n",
       "  <polygon points=\"0.0,0.0 120.0,0.0 120.0,29.030629010473877 0.0,29.030629010473877\" style=\"fill:#ECB172A0;stroke-width:0\"/>\n",
       "\n",
       "  <!-- Text -->\n",
       "  <text x=\"60.000000\" y=\"49.030629\" font-size=\"1.0rem\" font-weight=\"100\" text-anchor=\"middle\" >50</text>\n",
       "  <text x=\"140.000000\" y=\"14.515315\" font-size=\"1.0rem\" font-weight=\"100\" text-anchor=\"middle\" transform=\"rotate(0,140.000000,14.515315)\">1</text>\n",
       "</svg>\n",
       "        </td>\n",
       "    </tr>\n",
       "</table>"
      ],
      "text/plain": [
       "dask.array<getitem, shape=(50,), dtype=float64, chunksize=(50,), chunktype=numpy.ndarray>"
      ]
     },
     "execution_count": 30,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "a[:50, 200]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 31,
   "id": "ec7db820-07b0-4200-9dcc-5b1074547ac0",
   "metadata": {
    "editable": true,
    "slideshow": {
     "slide_type": "fragment"
    },
    "tags": []
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "array([-1.43243322e+00, -1.16493758e+00, -1.87967293e+00,  1.47886369e-01,\n",
       "        1.28971173e+00, -5.20397767e-01,  7.38853770e-01, -6.29040623e-01,\n",
       "        1.08457857e+00, -1.65313786e+00,  4.96115440e-02,  1.14667021e+00,\n",
       "       -1.84882693e-01, -4.90448809e-01, -1.98556060e-01, -4.59505020e-01,\n",
       "       -1.23344705e+00, -4.75947121e-01, -1.08411314e+00,  1.10891156e+00,\n",
       "       -3.80255131e-01, -2.05722915e-01,  1.76842336e+00,  2.14955832e-01,\n",
       "        4.06571959e-01,  1.07101227e+00,  1.70421791e+00, -2.57276521e+00,\n",
       "        2.11179670e+00,  9.94971691e-04,  6.76553884e-02,  1.02370386e+00,\n",
       "       -1.28329307e+00,  2.35121424e-01, -1.23902863e-01,  7.26158694e-01,\n",
       "        1.21912210e-01,  1.06211882e+00, -5.90280767e-01,  8.41471468e-01,\n",
       "       -4.44307566e-01, -2.65149637e-01,  1.28920793e+00,  1.83402709e+00,\n",
       "        1.66868151e+00,  3.36507699e-01, -1.54065640e-01,  1.04694489e+00,\n",
       "        4.46484496e-01, -9.05282945e-01])"
      ]
     },
     "execution_count": 31,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "a[:50, 100].compute()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 32,
   "id": "18d120ec-11cc-4efd-a0fd-aa258fc13b68",
   "metadata": {
    "editable": true,
    "slideshow": {
     "slide_type": ""
    },
    "tags": []
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "0.00214943349359884"
      ]
     },
     "execution_count": 32,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "a.mean().compute()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ff1a2239-9d9f-49cb-af12-fd769a6d2068",
   "metadata": {
    "editable": true,
    "slideshow": {
     "slide_type": "subslide"
    },
    "tags": []
   },
   "source": [
    "## Dask Distributed with Slurm"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7703556e-2736-415c-872f-634e9b89c7d2",
   "metadata": {
    "editable": true,
    "slideshow": {
     "slide_type": "fragment"
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "from dask.distributed import Client\n",
    "from dask_jobqueue import SLURMCluster"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "830e03d1-3e4d-4efa-8496-bfc8bf91a755",
   "metadata": {
    "editable": true,
    "slideshow": {
     "slide_type": "subslide"
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "# Define single unit of the Dask Distributed \"Cluster\"\n",
    "cluster = SLURMCluster(queue='admintest', cores=1, memory=\"20GB\")\n",
    "\n",
    "# Scale up the cluster to have 10 members\n",
    "cluster.scale(12)\n",
    "\n",
    "# Initialize the \"client\" so that the script is connected to the Cluster\n",
    "client = Client(cluster)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "47f0f3aa-182f-4147-a582-365b4c3f753f",
   "metadata": {
    "editable": true,
    "slideshow": {
     "slide_type": "subslide"
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "client"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d5c89217-34d8-4730-b695-8954d5267718",
   "metadata": {
    "editable": true,
    "slideshow": {
     "slide_type": "subslide"
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "data = np.random.normal(size=200000000).reshape(40000, 5000)\n",
    "a = da.from_array(data, chunks=(2000, 1000))\n",
    "a"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "26c7639b-16c6-45dd-a22f-50761c305df5",
   "metadata": {
    "editable": true,
    "slideshow": {
     "slide_type": "fragment"
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "a.std().compute()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9047066a-fd11-497e-bcab-577d5e86d076",
   "metadata": {
    "editable": true,
    "slideshow": {
     "slide_type": "subslide"
    },
    "tags": []
   },
   "source": [
    "## Example 3: NYC Taxi Data\n",
    "\n",
    "- Collected data from all taxi and ride-share trips\n",
    "- Very large data sets, too big to work with all at once on a single computer\n",
    "- Let's use `dask` to explore some facets of the data\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "990d7f6a-683d-4750-a8a4-83613fd1e89a",
   "metadata": {
    "editable": true,
    "slideshow": {
     "slide_type": "subslide"
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "yellow_cab = glob.glob('/home/tl397/ycrc/workshops/taxi/yellow_tripdata_2022-*parquet')\n",
    "ride_share = glob.glob('/home/tl397/ycrc/workshops/taxi/fhvhv_tripdata_2022-*parquet')\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2ea8c9fe-d6fd-49ec-8713-461958252202",
   "metadata": {
    "editable": true,
    "slideshow": {
     "slide_type": "fragment"
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "yc = dd.read_parquet(yellow_cab)\n",
    "rs = dd.read_parquet(ride_share) \n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6de5586b-f596-4b33-bc62-d28c48c4672d",
   "metadata": {
    "editable": true,
    "slideshow": {
     "slide_type": "fragment"
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "yc = yc[(yc.fare_amount > 0)]\n",
    "rs = rs[(rs.base_passenger_fare > 0)]\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "cad4ca6f-e27b-475c-80bd-ddd0a42076f9",
   "metadata": {
    "editable": true,
    "slideshow": {
     "slide_type": "subslide"
    },
    "tags": []
   },
   "source": [
    "### Question: Do people tip cabs or Ubers/Lyfts better?"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d01d3ba3-3b09-42cb-a61c-9129736b6555",
   "metadata": {
    "editable": true,
    "slideshow": {
     "slide_type": "fragment"
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "h_yc, bins = da.histogram(np.divide(yc.tip_amount, yc.fare_amount), bins=200, range=[0.01, 2])\n",
    "h_rs, bins = da.histogram(np.divide(rs.tips, rs.base_passenger_fare), bins=200, range=[0.01, 2])\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "99733e61-3014-4762-a73c-32170ea277c4",
   "metadata": {
    "editable": true,
    "slideshow": {
     "slide_type": "subslide"
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "plt.subplots(1,1)\n",
    "plt.stairs(h_yc, bins, label=\"yellow cab\")\n",
    "plt.stairs(h_rs, bins, label=\"uber/lyft\")\n",
    "\n",
    "plt.yscale('log');\n",
    "plt.ylabel('Rides');\n",
    "plt.xlabel('Tip percentage (%)');\n",
    "plt.legend();"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "64dc76c8-0eeb-491c-9193-24831dda0b29",
   "metadata": {
    "editable": true,
    "slideshow": {
     "slide_type": "subslide"
    },
    "tags": []
   },
   "source": [
    "### Mean tip percentage"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "09ddcdc4-c2ac-4c64-9228-8fdea7225640",
   "metadata": {
    "editable": true,
    "scrolled": true,
    "slideshow": {
     "slide_type": "fragment"
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "print(f\"Yellow Cab: {100*yc.tip_amount.divide(yc.fare_amount).mean().compute():.2f}%\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1aec4f3f-9066-442d-b766-83ad896b2497",
   "metadata": {
    "editable": true,
    "slideshow": {
     "slide_type": "fragment"
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "print(f\"Ride-share: {100*rs.tips.divide(rs.base_passenger_fare).mean().compute():.2f}%\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1d3ed5ed-7644-4a6c-9297-7b197cf22e16",
   "metadata": {
    "editable": true,
    "slideshow": {
     "slide_type": "subslide"
    },
    "tags": []
   },
   "source": [
    "## Key Take-aways\n",
    "1. Dask is able to orchestrate lots of parallel workers, either locally or across the cluster\n",
    "2. It's easier to not tip when it's on an app?\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.12.8"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}