benchmark_gcm-filters.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "c3c9e2b6-4e85-45a8-8f55-342a54a7a43f",
   "metadata": {},
   "source": [
    "# GCM Filters Scaling Benchmark\n",
    "\n",
    "Run on Casper full node (36 cores)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "789dda38-3ccb-48ed-8658-a741771cb5a3",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "'0.18.2'"
      ]
     },
     "execution_count": 1,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "import xarray as xr\n",
    "xr.__version__"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "1afa2aaa-a193-4f1a-8809-68dc4bafe519",
   "metadata": {},
   "outputs": [],
   "source": [
    "#from coiled import performance_report\n",
    "from dask.distributed import performance_report"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "689867dc-a6e4-4ee6-b5d4-6e1a193183ce",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "'2021.06.2'"
      ]
     },
     "execution_count": 3,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "import dask\n",
    "dask.__version__"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "276cf7b3-d432-45a4-8cfa-b8a70310994c",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "'1.18.2'"
      ]
     },
     "execution_count": 4,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "import numpy as np\n",
    "np.__version__"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "e24e9b99-0af2-478b-b3a3-a7a1dcc43c4c",
   "metadata": {},
   "outputs": [
    {
     "ename": "ImportError",
     "evalue": "CuPy is not correctly installed.\n\nIf you are using wheel distribution (cupy-cudaXX), make sure that the version of CuPy you installed matches with the version of CUDA on your host.\nAlso, confirm that only one CuPy package is installed:\n  $ pip freeze\n\nIf you are building CuPy from source, please check your environment, uninstall CuPy and reinstall it with:\n  $ pip install cupy --no-cache-dir -vvvv\n\nCheck the Installation Guide for details:\n  https://docs-cupy.chainer.org/en/latest/install.html\n\noriginal error: libcublas.so.10: cannot open shared object file: No such file or directory",
     "output_type": "error",
     "traceback": [
      "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
      "\u001b[0;31mImportError\u001b[0m                               Traceback (most recent call last)",
      "\u001b[0;32m/glade/work/rpa/my_npl_clone/lib/python3.7/site-packages/cupy/__init__.py\u001b[0m in \u001b[0;36m<module>\u001b[0;34m\u001b[0m\n\u001b[1;32m     20\u001b[0m                                 message='can\\'t resolve package from __spec__')\n\u001b[0;32m---> 21\u001b[0;31m         \u001b[0;32mfrom\u001b[0m \u001b[0mcupy\u001b[0m \u001b[0;32mimport\u001b[0m \u001b[0mcore\u001b[0m  \u001b[0;31m# NOQA\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m     22\u001b[0m \u001b[0;32mexcept\u001b[0m \u001b[0mImportError\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;32m/glade/work/rpa/my_npl_clone/lib/python3.7/site-packages/cupy/core/__init__.py\u001b[0m in \u001b[0;36m<module>\u001b[0;34m\u001b[0m\n\u001b[0;32m----> 1\u001b[0;31m \u001b[0;32mfrom\u001b[0m \u001b[0mcupy\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mcore\u001b[0m \u001b[0;32mimport\u001b[0m \u001b[0mcore\u001b[0m  \u001b[0;31m# NOQA\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m      2\u001b[0m \u001b[0;32mfrom\u001b[0m \u001b[0mcupy\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mcore\u001b[0m \u001b[0;32mimport\u001b[0m \u001b[0minternal\u001b[0m  \u001b[0;31m# NOQA\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m      3\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;31mImportError\u001b[0m: libcublas.so.10: cannot open shared object file: No such file or directory",
      "\nDuring handling of the above exception, another exception occurred:\n",
      "\u001b[0;31mImportError\u001b[0m                               Traceback (most recent call last)",
      "\u001b[0;32m<ipython-input-6-9ffcf30d9aea>\u001b[0m in \u001b[0;36m<module>\u001b[0;34m\u001b[0m\n\u001b[0;32m----> 1\u001b[0;31m \u001b[0;32mimport\u001b[0m \u001b[0mcupy\u001b[0m \u001b[0;32mas\u001b[0m \u001b[0mcp\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m      2\u001b[0m \u001b[0mcp\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0m__version__\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;32m/glade/work/rpa/my_npl_clone/lib/python3.7/site-packages/cupy/__init__.py\u001b[0m in \u001b[0;36m<module>\u001b[0;34m\u001b[0m\n\u001b[1;32m     40\u001b[0m original error: {}'''.format(exc_info[1]))  # NOQA\n\u001b[1;32m     41\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m---> 42\u001b[0;31m     \u001b[0msix\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mreraise\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mImportError\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mImportError\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mmsg\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mexc_info\u001b[0m\u001b[0;34m[\u001b[0m\u001b[0;36m2\u001b[0m\u001b[0;34m]\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m     43\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m     44\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;32m/glade/work/rpa/my_npl_clone/lib/python3.7/site-packages/six.py\u001b[0m in \u001b[0;36mreraise\u001b[0;34m(tp, value, tb)\u001b[0m\n\u001b[1;32m    700\u001b[0m                 \u001b[0mvalue\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mtp\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    701\u001b[0m             \u001b[0;32mif\u001b[0m \u001b[0mvalue\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0m__traceback__\u001b[0m \u001b[0;32mis\u001b[0m \u001b[0;32mnot\u001b[0m \u001b[0mtb\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m--> 702\u001b[0;31m                 \u001b[0;32mraise\u001b[0m \u001b[0mvalue\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mwith_traceback\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mtb\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m    703\u001b[0m             \u001b[0;32mraise\u001b[0m \u001b[0mvalue\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    704\u001b[0m         \u001b[0;32mfinally\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;32m/glade/work/rpa/my_npl_clone/lib/python3.7/site-packages/cupy/__init__.py\u001b[0m in \u001b[0;36m<module>\u001b[0;34m\u001b[0m\n\u001b[1;32m     19\u001b[0m         warnings.filterwarnings('ignore', category=ImportWarning,\n\u001b[1;32m     20\u001b[0m                                 message='can\\'t resolve package from __spec__')\n\u001b[0;32m---> 21\u001b[0;31m         \u001b[0;32mfrom\u001b[0m \u001b[0mcupy\u001b[0m \u001b[0;32mimport\u001b[0m \u001b[0mcore\u001b[0m  \u001b[0;31m# NOQA\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m     22\u001b[0m \u001b[0;32mexcept\u001b[0m \u001b[0mImportError\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m     23\u001b[0m     \u001b[0;31m# core is a c-extension module.\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;32m/glade/work/rpa/my_npl_clone/lib/python3.7/site-packages/cupy/core/__init__.py\u001b[0m in \u001b[0;36m<module>\u001b[0;34m\u001b[0m\n\u001b[0;32m----> 1\u001b[0;31m \u001b[0;32mfrom\u001b[0m \u001b[0mcupy\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mcore\u001b[0m \u001b[0;32mimport\u001b[0m \u001b[0mcore\u001b[0m  \u001b[0;31m# NOQA\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m      2\u001b[0m \u001b[0;32mfrom\u001b[0m \u001b[0mcupy\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mcore\u001b[0m \u001b[0;32mimport\u001b[0m \u001b[0minternal\u001b[0m  \u001b[0;31m# NOQA\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m      3\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m      4\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m      5\u001b[0m \u001b[0;31m# import class and function\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;31mImportError\u001b[0m: CuPy is not correctly installed.\n\nIf you are using wheel distribution (cupy-cudaXX), make sure that the version of CuPy you installed matches with the version of CUDA on your host.\nAlso, confirm that only one CuPy package is installed:\n  $ pip freeze\n\nIf you are building CuPy from source, please check your environment, uninstall CuPy and reinstall it with:\n  $ pip install cupy --no-cache-dir -vvvv\n\nCheck the Installation Guide for details:\n  https://docs-cupy.chainer.org/en/latest/install.html\n\noriginal error: libcublas.so.10: cannot open shared object file: No such file or directory"
     ]
    }
   ],
   "source": [
    "import cupy as cp\n",
    "cp.__version__"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "d96eb689-dc09-4f44-ab13-074a803724ad",
   "metadata": {},
   "outputs": [],
   "source": [
    "import dask.array as dsa"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "e0a7d073-c150-427b-b70d-849f1749506d",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "'0.1'"
      ]
     },
     "execution_count": 6,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "import gcm_filters as gf\n",
    "gf.__version__"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "6d8267db-a143-4381-950a-be5bf506de1c",
   "metadata": {},
   "outputs": [],
   "source": [
    "from contextlib import contextmanager\n",
    "import time\n",
    "import pandas as pd\n",
    "\n",
    "class DiagnosticTimer:\n",
    "    def __init__(self):\n",
    "        self.diagnostics = []\n",
    "\n",
    "    @contextmanager\n",
    "    def time(self, **kwargs):\n",
    "        tic = time.time()\n",
    "        yield\n",
    "        toc = time.time()\n",
    "        kwargs[\"runtime\"] = toc - tic\n",
    "        self.diagnostics.append(kwargs)\n",
    "\n",
    "    @property\n",
    "    def df(self):\n",
    "        return pd.DataFrame(self.diagnostics)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "16a92dfc-fc2e-4894-b1b1-895cc67af353",
   "metadata": {},
   "outputs": [],
   "source": [
    "def make_data(shape, chunks, gpu=False):\n",
    "    nt, ny, nx = shape\n",
    "    da = xr.DataArray(dsa.random.random(shape, chunks=chunks), dims=['time', 'y', 'x'])\n",
    "    mask_data = dsa.ones((ny, nx))\n",
    "    mask_data[(ny // 4):(3 * ny // 4), (nx // 4):(3 * nx // 4)] = 0\n",
    "    wet_mask = xr.DataArray(mask_data, dims=['y', 'x'])\n",
    "    \n",
    "    da_masked = da.where(wet_mask)\n",
    "    \n",
    "    if gpu:\n",
    "        raise NotImplementedError(\"Can't get cupy working :(\")\n",
    "\n",
    "    filter = gf.Filter(\n",
    "        filter_scale=4,\n",
    "        dx_min=1,\n",
    "        filter_shape=gf.FilterShape.TAPER,\n",
    "        grid_type=gf.GridType.REGULAR_WITH_LAND,\n",
    "        grid_vars={'wet_mask': wet_mask}\n",
    "    )\n",
    "    filter\n",
    "\n",
    "    da_filtered = filter.apply(da_masked, dims=['y', 'x'])\n",
    "    return da_masked, da_filtered"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "8c594e24-9cbd-4517-a4f1-f992c8108675",
   "metadata": {},
   "outputs": [],
   "source": [
    "from dask.distributed import Client, LocalCluster"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "id": "21f48010-763f-41f3-987d-dafd38d24054",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/glade/work/rpa/my_npl_clone/lib/python3.7/site-packages/distributed/node.py:164: UserWarning: Port 8787 is already in use.\n",
      "Perhaps you already have a cluster running?\n",
      "Hosting the HTTP server on port 33998 instead\n",
      "  expected, actual\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "         operation  ncores  nworkers              shape            chunks  \\\n",
      "0  unfiltered_mean       3         1  (360, 1024, 1024)  (10, 1024, 1024)   \n",
      "1    filtered_mean       3         1  (360, 1024, 1024)  (10, 1024, 1024)   \n",
      "\n",
      "       nbytes    dtype    runtime  \n",
      "0  3019898880  float64   1.928494  \n",
      "1  3019898880  float64  86.083285  \n",
      "         operation  ncores  nworkers              shape            chunks  \\\n",
      "2  unfiltered_mean       9         3  (360, 1024, 1024)  (10, 1024, 1024)   \n",
      "3    filtered_mean       9         3  (360, 1024, 1024)  (10, 1024, 1024)   \n",
      "\n",
      "       nbytes    dtype    runtime  \n",
      "2  3019898880  float64   0.806149  \n",
      "3  3019898880  float64  31.352030  \n",
      "         operation  ncores  nworkers              shape            chunks  \\\n",
      "4  unfiltered_mean      18         6  (360, 1024, 1024)  (10, 1024, 1024)   \n",
      "5    filtered_mean      18         6  (360, 1024, 1024)  (10, 1024, 1024)   \n",
      "\n",
      "       nbytes    dtype    runtime  \n",
      "4  3019898880  float64   0.563238  \n",
      "5  3019898880  float64  19.720357  \n",
      "         operation  ncores  nworkers              shape            chunks  \\\n",
      "6  unfiltered_mean      27         9  (360, 1024, 1024)  (10, 1024, 1024)   \n",
      "7    filtered_mean      27         9  (360, 1024, 1024)  (10, 1024, 1024)   \n",
      "\n",
      "       nbytes    dtype    runtime  \n",
      "6  3019898880  float64   0.441960  \n",
      "7  3019898880  float64  18.742256  \n",
      "         operation  ncores  nworkers              shape            chunks  \\\n",
      "8  unfiltered_mean      36        12  (360, 1024, 1024)  (10, 1024, 1024)   \n",
      "9    filtered_mean      36        12  (360, 1024, 1024)  (10, 1024, 1024)   \n",
      "\n",
      "       nbytes    dtype    runtime  \n",
      "8  3019898880  float64   0.438341  \n",
      "9  3019898880  float64  15.171387  \n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/glade/work/rpa/my_npl_clone/lib/python3.7/site-packages/distributed/node.py:164: UserWarning: Port 8787 is already in use.\n",
      "Perhaps you already have a cluster running?\n",
      "Hosting the HTTP server on port 39942 instead\n",
      "  expected, actual\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "          operation  ncores  nworkers              shape            chunks  \\\n",
      "10  unfiltered_mean       9         1  (360, 1024, 1024)  (10, 1024, 1024)   \n",
      "11    filtered_mean       9         1  (360, 1024, 1024)  (10, 1024, 1024)   \n",
      "\n",
      "        nbytes    dtype    runtime  \n",
      "10  3019898880  float64   0.826083  \n",
      "11  3019898880  float64  33.627329  \n",
      "          operation  ncores  nworkers              shape            chunks  \\\n",
      "12  unfiltered_mean       9         1  (360, 1024, 1024)  (10, 1024, 1024)   \n",
      "13    filtered_mean       9         1  (360, 1024, 1024)  (10, 1024, 1024)   \n",
      "\n",
      "        nbytes    dtype    runtime  \n",
      "12  3019898880  float64   0.701733  \n",
      "13  3019898880  float64  33.094337  \n",
      "          operation  ncores  nworkers              shape            chunks  \\\n",
      "14  unfiltered_mean      18         2  (360, 1024, 1024)  (10, 1024, 1024)   \n",
      "15    filtered_mean      18         2  (360, 1024, 1024)  (10, 1024, 1024)   \n",
      "\n",
      "        nbytes    dtype    runtime  \n",
      "14  3019898880  float64   0.491663  \n",
      "15  3019898880  float64  19.829080  \n",
      "          operation  ncores  nworkers              shape            chunks  \\\n",
      "16  unfiltered_mean      27         3  (360, 1024, 1024)  (10, 1024, 1024)   \n",
      "17    filtered_mean      27         3  (360, 1024, 1024)  (10, 1024, 1024)   \n",
      "\n",
      "        nbytes    dtype    runtime  \n",
      "16  3019898880  float64   0.422401  \n",
      "17  3019898880  float64  19.782855  \n",
      "          operation  ncores  nworkers              shape            chunks  \\\n",
      "18  unfiltered_mean      36         4  (360, 1024, 1024)  (10, 1024, 1024)   \n",
      "19    filtered_mean      36         4  (360, 1024, 1024)  (10, 1024, 1024)   \n",
      "\n",
      "        nbytes    dtype    runtime  \n",
      "18  3019898880  float64   0.421241  \n",
      "19  3019898880  float64  15.231666  \n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/glade/work/rpa/my_npl_clone/lib/python3.7/site-packages/distributed/node.py:164: UserWarning: Port 8787 is already in use.\n",
      "Perhaps you already have a cluster running?\n",
      "Hosting the HTTP server on port 37744 instead\n",
      "  expected, actual\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "          operation  ncores  nworkers              shape            chunks  \\\n",
      "20  unfiltered_mean      36         1  (360, 1024, 1024)  (10, 1024, 1024)   \n",
      "21    filtered_mean      36         1  (360, 1024, 1024)  (10, 1024, 1024)   \n",
      "\n",
      "        nbytes    dtype    runtime  \n",
      "20  3019898880  float64   0.527528  \n",
      "21  3019898880  float64  18.852935  \n",
      "          operation  ncores  nworkers              shape            chunks  \\\n",
      "22  unfiltered_mean      36         1  (360, 1024, 1024)  (10, 1024, 1024)   \n",
      "23    filtered_mean      36         1  (360, 1024, 1024)  (10, 1024, 1024)   \n",
      "\n",
      "        nbytes    dtype    runtime  \n",
      "22  3019898880  float64   0.420205  \n",
      "23  3019898880  float64  18.078998  \n"
     ]
    }
   ],
   "source": [
    "# strong scaling - problem size stays the same\n",
    "\n",
    "shape = 360, 1024, 1024\n",
    "chunks = (10,) + shape[1:]\n",
    "\n",
    "unfiltered, filtered = make_data(shape, chunks)\n",
    "\n",
    "diag_timer_strong = DiagnosticTimer()\n",
    "\n",
    "for threads_per_worker in [3, 9, 36]:\n",
    "    max_workers = 36 // threads_per_worker\n",
    "    worker_step = max(max_workers // 4, 1)\n",
    "    cluster = LocalCluster(threads_per_worker=threads_per_worker, n_workers=1)\n",
    "    client = Client(cluster)\n",
    "    for nworkers in [1] + list(range(worker_step, max_workers + 1, worker_step)):\n",
    "        cluster.scale(nworkers)\n",
    "        client.wait_for_workers(nworkers)\n",
    "        assert len(client.ncores()) == nworkers\n",
    "        ncores = sum(client.ncores().values())\n",
    "        details = dict(ncores=ncores, nworkers=nworkers, shape=shape, chunks=chunks,\n",
    "                       nbytes=filtered.data.nbytes, dtype=str(filtered.dtype))\n",
    "        with diag_timer_strong.time(operation='unfiltered_mean', **details):\n",
    "            unfiltered.data.mean().compute()\n",
    "        with diag_timer_strong.time(operation='filtered_mean', **details):\n",
    "            filtered.data.mean().compute()\n",
    "        print(diag_timer_strong.df.iloc[-2:])\n",
    "    client.close()\n",
    "    cluster.close()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "id": "836dc719-85ad-4521-a915-365c710145e7",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/glade/work/rpa/my_npl_clone/lib/python3.7/site-packages/distributed/node.py:164: UserWarning: Port 8787 is already in use.\n",
      "Perhaps you already have a cluster running?\n",
      "Hosting the HTTP server on port 41355 instead\n",
      "  expected, actual\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "         operation  ncores  nworkers              shape            chunks  \\\n",
      "0  unfiltered_mean       3         1  (120, 1024, 1024)  (10, 1024, 1024)   \n",
      "1    filtered_mean       3         1  (120, 1024, 1024)  (10, 1024, 1024)   \n",
      "\n",
      "       nbytes    dtype    runtime  \n",
      "0  1006632960  float64   0.755714  \n",
      "1  1006632960  float64  29.135221  \n",
      "         operation  ncores  nworkers              shape            chunks  \\\n",
      "2  unfiltered_mean       9         3  (360, 1024, 1024)  (10, 1024, 1024)   \n",
      "3    filtered_mean       9         3  (360, 1024, 1024)  (10, 1024, 1024)   \n",
      "\n",
      "       nbytes    dtype    runtime  \n",
      "2  3019898880  float64   0.811971  \n",
      "3  3019898880  float64  30.836951  \n",
      "         operation  ncores  nworkers              shape            chunks  \\\n",
      "4  unfiltered_mean      18         6  (720, 1024, 1024)  (10, 1024, 1024)   \n",
      "5    filtered_mean      18         6  (720, 1024, 1024)  (10, 1024, 1024)   \n",
      "\n",
      "       nbytes    dtype    runtime  \n",
      "4  6039797760  float64   0.823065  \n",
      "5  6039797760  float64  35.334654  \n",
      "         operation  ncores  nworkers               shape            chunks  \\\n",
      "6  unfiltered_mean      27         9  (1080, 1024, 1024)  (10, 1024, 1024)   \n",
      "7    filtered_mean      27         9  (1080, 1024, 1024)  (10, 1024, 1024)   \n",
      "\n",
      "       nbytes    dtype    runtime  \n",
      "6  9059696640  float64   1.049813  \n",
      "7  9059696640  float64  45.319927  \n",
      "         operation  ncores  nworkers               shape            chunks  \\\n",
      "8  unfiltered_mean      36        12  (1440, 1024, 1024)  (10, 1024, 1024)   \n",
      "9    filtered_mean      36        12  (1440, 1024, 1024)  (10, 1024, 1024)   \n",
      "\n",
      "        nbytes    dtype    runtime  \n",
      "8  12079595520  float64   1.123856  \n",
      "9  12079595520  float64  58.703441  \n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/glade/work/rpa/my_npl_clone/lib/python3.7/site-packages/distributed/node.py:164: UserWarning: Port 8787 is already in use.\n",
      "Perhaps you already have a cluster running?\n",
      "Hosting the HTTP server on port 34211 instead\n",
      "  expected, actual\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "          operation  ncores  nworkers              shape            chunks  \\\n",
      "10  unfiltered_mean       9         1  (360, 1024, 1024)  (10, 1024, 1024)   \n",
      "11    filtered_mean       9         1  (360, 1024, 1024)  (10, 1024, 1024)   \n",
      "\n",
      "        nbytes    dtype    runtime  \n",
      "10  3019898880  float64   0.842147  \n",
      "11  3019898880  float64  33.507577  \n",
      "          operation  ncores  nworkers              shape            chunks  \\\n",
      "12  unfiltered_mean       9         1  (360, 1024, 1024)  (10, 1024, 1024)   \n",
      "13    filtered_mean       9         1  (360, 1024, 1024)  (10, 1024, 1024)   \n",
      "\n",
      "        nbytes    dtype    runtime  \n",
      "12  3019898880  float64   0.714265  \n",
      "13  3019898880  float64  32.872067  \n",
      "          operation  ncores  nworkers              shape            chunks  \\\n",
      "14  unfiltered_mean      18         2  (720, 1024, 1024)  (10, 1024, 1024)   \n",
      "15    filtered_mean      18         2  (720, 1024, 1024)  (10, 1024, 1024)   \n",
      "\n",
      "        nbytes    dtype    runtime  \n",
      "14  6039797760  float64   0.857916  \n",
      "15  6039797760  float64  38.851158  \n",
      "          operation  ncores  nworkers               shape            chunks  \\\n",
      "16  unfiltered_mean      27         3  (1080, 1024, 1024)  (10, 1024, 1024)   \n",
      "17    filtered_mean      27         3  (1080, 1024, 1024)  (10, 1024, 1024)   \n",
      "\n",
      "        nbytes    dtype    runtime  \n",
      "16  9059696640  float64   0.891048  \n",
      "17  9059696640  float64  46.904443  \n",
      "          operation  ncores  nworkers               shape            chunks  \\\n",
      "18  unfiltered_mean      36         4  (1440, 1024, 1024)  (10, 1024, 1024)   \n",
      "19    filtered_mean      36         4  (1440, 1024, 1024)  (10, 1024, 1024)   \n",
      "\n",
      "         nbytes    dtype    runtime  \n",
      "18  12079595520  float64   1.026639  \n",
      "19  12079595520  float64  60.178699  \n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/glade/work/rpa/my_npl_clone/lib/python3.7/site-packages/distributed/node.py:164: UserWarning: Port 8787 is already in use.\n",
      "Perhaps you already have a cluster running?\n",
      "Hosting the HTTP server on port 45430 instead\n",
      "  expected, actual\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "          operation  ncores  nworkers               shape            chunks  \\\n",
      "20  unfiltered_mean      36         1  (1440, 1024, 1024)  (10, 1024, 1024)   \n",
      "21    filtered_mean      36         1  (1440, 1024, 1024)  (10, 1024, 1024)   \n",
      "\n",
      "         nbytes    dtype    runtime  \n",
      "20  12079595520  float64   1.381825  \n",
      "21  12079595520  float64  72.642258  \n",
      "          operation  ncores  nworkers               shape            chunks  \\\n",
      "22  unfiltered_mean      36         1  (1440, 1024, 1024)  (10, 1024, 1024)   \n",
      "23    filtered_mean      36         1  (1440, 1024, 1024)  (10, 1024, 1024)   \n",
      "\n",
      "         nbytes    dtype    runtime  \n",
      "22  12079595520  float64   1.128636  \n",
      "23  12079595520  float64  72.032173  \n"
     ]
    }
   ],
   "source": [
    "# weak scaling - problem size scales with number of cores\n",
    "\n",
    "diag_timer_weak = DiagnosticTimer()\n",
    "\n",
    "for threads_per_worker in [3, 9, 36]:\n",
    "    max_workers = 36 // threads_per_worker\n",
    "    worker_step = max(max_workers // 4, 1)\n",
    "    cluster = LocalCluster(threads_per_worker=threads_per_worker, n_workers=1)\n",
    "    client = Client(cluster)\n",
    "    for nworkers in [1] + list(range(worker_step, max_workers + 1, worker_step)):\n",
    "        cluster.scale(nworkers)\n",
    "        client.wait_for_workers(nworkers)\n",
    "        assert len(client.ncores()) == nworkers\n",
    "        ncores = sum(client.ncores().values())\n",
    "        \n",
    "        shape = ncores * 10 * 4, 1024, 1024\n",
    "        chunks = (10,) + shape[1:]\n",
    "        unfiltered, filtered = make_data(shape, chunks)\n",
    "\n",
    "        details = dict(ncores=ncores, nworkers=nworkers, shape=shape, chunks=chunks,\n",
    "                       nbytes=filtered.data.nbytes, dtype=str(filtered.dtype))\n",
    "        with diag_timer_weak.time(operation='unfiltered_mean', **details):\n",
    "            unfiltered.data.mean().compute()\n",
    "        with diag_timer_weak.time(operation='filtered_mean', **details):\n",
    "            filtered.data.mean().compute()\n",
    "        print(diag_timer_weak.df.iloc[-2:])\n",
    "    client.close()\n",
    "    cluster.close()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "id": "d4e940ed-8242-4a97-80d3-4df90cdf25b0",
   "metadata": {},
   "outputs": [],
   "source": [
    "from datetime import datetime\n",
    "\n",
    "now = datetime.now().isoformat()[:19]\n",
    "diag_timer_weak.df.to_csv(f'data/scaling_weak_cpu_{now}.csv', index=False)\n",
    "diag_timer_strong.df.to_csv(f'data/scaling_strong_cpu_{now}.csv', index=False)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "d94a0872-19e7-417f-b1fe-2635977c6d29",
   "metadata": {},
   "outputs": [],
   "source": [
    "# performance report\n",
    "\n",
    "threads_per_worker = 9\n",
    "nworkers = 4\n",
    "\n",
    "cluster = LocalCluster(threads_per_worker=threads_per_worker, n_workers=1)\n",
    "client = Client(cluster)\n",
    "cluster.scale(nworkers)\n",
    "client.wait_for_workers(nworkers)\n",
    "assert len(client.ncores()) == nworkers\n",
    "ncores = sum(client.ncores().values())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "id": "48f0f6e5-d615-43a3-b6d5-f8d49924da34",
   "metadata": {},
   "outputs": [],
   "source": [
    "shape = ncores * 10 * 4, 1024, 1024\n",
    "chunks = (10,) + shape[1:]\n",
    "unfiltered, filtered = make_data(shape, chunks)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "id": "fc29efb8-4602-4e93-9e48-66879e5e263d",
   "metadata": {},
   "outputs": [],
   "source": [
    "with performance_report(\"performance_reports/unfiltered_mean_36_cores_4_workers.html\"):\n",
    "    unfiltered.data.mean().compute()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "id": "2597eb21-3c4b-4ac3-bfd6-004068bdde34",
   "metadata": {},
   "outputs": [],
   "source": [
    "with performance_report(\"performance_reports/filtered_mean_36_cores_4_workers.html\"):\n",
    "    filtered.data.mean().compute()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9a02b611-bb6c-4d5d-ba68-2a91f7f4a0c0",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "NPL3 (my_npl_clone)",
   "language": "python",
   "name": "npl3-my_npl_clone"
  },
  "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.7.5"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}