Parallelism with Joblib.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# A tour of python profiling and optimization (Part II)\n",
    "\n",
    "*Antonino Ingargiola*, EuroScipy 2018\n",
    "\n",
    "> premature optimization is the root of all evil (or at least most of it) in programming. \n",
    "\n",
    "                                                         - Donald Knuth\n",
    "\n",
    "# Parallelism with Joblib\n",
    "\n",
    "> [Joblib](https://joblib.readthedocs.io/) provides a *incredibly simple API* for both \n",
    "> **thread-based** a **process-based** parallelism (it also offer memoization).\n",
    "\n",
    "\n",
    "*See also Part I: [Optimizing python code](Optimizing%20python%20code.ipynb).*\n",
    "\n",
    "**Binder Warning**: If you are running this on *MyBinder.org* you may not see any\n",
    "speed-up because each user is limited to 1 CPU."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## OU process function to be parallelized\n",
    "\n",
    "*Continuing the Ornestein-Uhlenbeck process example from [Optimizing python code](Optimizing%20python%20code.ipynb).*\n",
    "\n",
    "We define both Cython and Numba functions to be parallelized:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import time, os\n",
    "import numpy as np\n",
    "import numba\n",
    "import math\n",
    "from randomgen import RandomGenerator, Xoroshiro128"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "%load_ext cython"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "%%cython\n",
    "\n",
    "import numpy as np\n",
    "cimport numpy as np\n",
    "from libc.math cimport exp, sqrt\n",
    "import cython\n",
    "\n",
    "@cython.boundscheck(False)\n",
    "def OU_process_cy3(int num_steps, *, double[:] N_norm, double dt=0.1, \n",
    "                   double x0=0, double tau=1, double sigma=2):\n",
    "    cdef Py_ssize_t i\n",
    "    cdef double relax, diffuse\n",
    "    cdef double[:] x = np.empty(num_steps + 1)\n",
    "    \n",
    "    with nogil:\n",
    "        relax = exp(-dt / tau)\n",
    "        diffuse = sigma * sqrt(1 - relax**2)\n",
    "        x[0] = x0\n",
    "        for i in range(num_steps):\n",
    "            x[i+1] = x[i] * relax + diffuse * N_norm[i]\n",
    "    return np.asfarray(x)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "@numba.jit(nogil=True)\n",
    "def OU_process_nb(num_steps, *, N_norm, δt=0.1, x0=0., τ=1., σ=2.):\n",
    "    relax = math.exp(-δt / τ)\n",
    "    diffuse = σ * math.sqrt(1 - relax**2)\n",
    "\n",
    "    x = np.empty(num_steps + 1)\n",
    "    x[0] = x0\n",
    "    for i in range(num_steps):\n",
    "        x[i+1] = x[i] * relax + diffuse * N_norm[i]\n",
    "    return x"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "> **NOTE**: Both functions release the GIL! Can you find where?"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Joblib examples\n",
    "\n",
    "*Reference: [Embarrassingly parallel for loops](https://joblib.readthedocs.io/en/latest/parallel.html)*"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "from joblib import Parallel, delayed"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Parallel sqrt\n",
    "\n",
    "Example from the docs:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "[0.0, 1.0, 2.0, 3.0, 4.0, 5.0, 6.0, 7.0, 8.0, 9.0]"
      ]
     },
     "execution_count": 6,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "[math.sqrt(i ** 2) for i in range(10)]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "[0.0, 1.0, 2.0, 3.0, 4.0, 5.0, 6.0, 7.0, 8.0, 9.0]"
      ]
     },
     "execution_count": 7,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "Parallel(n_jobs=2)(delayed(math.sqrt)(i ** 2) for i in range(10))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Busy sleep"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [],
   "source": [
    "def verbose_sleep(n):\n",
    "    print(f'PID: {os.getpid()} - busy sleeping', flush=True)\n",
    "    time.sleep(n)\n",
    "    print(f'PID: {os.getpid()} - done sleeping', flush=True)\n",
    "    return n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "PID: 54056 - busy sleeping\n",
      "PID: 54056 - done sleeping\n",
      "PID: 54056 - busy sleeping\n",
      "PID: 54056 - done sleeping\n",
      "PID: 54056 - busy sleeping\n",
      "PID: 54056 - done sleeping\n",
      "PID: 54056 - busy sleeping\n",
      "PID: 54056 - done sleeping\n",
      "20 s ± 0 ns per loop (mean ± std. dev. of 1 run, 1 loop each)\n"
     ]
    }
   ],
   "source": [
    "%%timeit -n1 -r1\n",
    "[verbose_sleep(5) for i in range(4)]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "5.4 s ± 0 ns per loop (mean ± std. dev. of 1 run, 1 loop each)\n"
     ]
    }
   ],
   "source": [
    "%%timeit -n1 -r1\n",
    "Parallel(n_jobs=4)(delayed(verbose_sleep)(5) for i in range(4))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Random Numbers Generations"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [],
   "source": [
    "num_steps = int(5e6)\n",
    "num_sim = 8\n",
    "\n",
    "rg = RandomGenerator(Xoroshiro128())"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "First the single-core version using `np.random`:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "1.55 s ± 144 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)\n"
     ]
    }
   ],
   "source": [
    "%%timeit\n",
    "np.random.randn(num_sim*num_steps)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Here we use thread-based parallelism (`prefer='threads'`):"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "1.31 s ± 82.5 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)\n"
     ]
    }
   ],
   "source": [
    "%%timeit\n",
    "Parallel(n_jobs=4, prefer='threads')(delayed(np.random.randn)(num_steps) \n",
    "                                     for i in range(num_sim))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now the comparison using `randomgen`:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "630 ms ± 52 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)\n"
     ]
    }
   ],
   "source": [
    "%%timeit\n",
    "rg.randn(num_sim*num_steps)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "580 ms ± 33.6 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)\n"
     ]
    }
   ],
   "source": [
    "%%timeit\n",
    "Parallel(n_jobs=4, prefer='threads')(delayed(rg.randn)(num_steps) \n",
    "                                     for i in range(num_sim))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Process-based parallelism, commented because too slow:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {},
   "outputs": [],
   "source": [
    "# %%timeit -n1 -r1\n",
    "# Parallel(n_jobs=4)(delayed(np.random.randn)(num_steps) for i in range(num_sim))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "metadata": {},
   "outputs": [],
   "source": [
    "# %%timeit -n1 -r1\n",
    "# Parallel(n_jobs=4)(delayed(rg.randn)(num_steps) for i in range(num_sim))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Parallelize the OU process\n",
    "\n",
    "In the previous notebook we have split the simulation in\n",
    "RNG and the deterministic computations. This is allows better\n",
    "optimization and to easily switch to the more performant `randomgen` \n",
    "instead of `numpy.random`.\n",
    "\n",
    "For parallelization purposes, here we define 2 wrapper function \n",
    "performing the full simulation (cython and numba versions):"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "metadata": {},
   "outputs": [],
   "source": [
    "def ou_process_cy_mp(num_steps, *, dt=0.1, x0=0, tau=1, sigma=2):\n",
    "    rg = RandomGenerator(Xoroshiro128())\n",
    "    N_norm = rg.randn(num_steps)\n",
    "    x = OU_process_cy3(num_steps, N_norm=N_norm, dt=dt, x0=x0, tau=tau, sigma=sigma)\n",
    "    return x\n",
    "\n",
    "def ou_process_nb_mp(num_steps, *, dt=0.1, x0=0, tau=1, sigma=2):\n",
    "    rg = RandomGenerator(Xoroshiro128())\n",
    "    N_norm = rg.randn(num_steps)\n",
    "    x = OU_process_nb(num_steps, N_norm=N_norm, δt=dt, x0=x0, τ=tau, σ=sigma)\n",
    "    return x"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Test calling the functions:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "[array([ 0.        , -1.35056857, -0.42634809, ..., -0.47090951,\n",
       "        -1.56315924, -2.29438242]),\n",
       " array([ 0.        ,  0.39933098, -0.50241349, ..., -1.04086923,\n",
       "        -1.45958082, -1.60680918]),\n",
       " array([ 0.        ,  0.96531688,  0.1634128 , ...,  0.72873354,\n",
       "        -1.08460205,  0.34895678]),\n",
       " array([ 0.        , -0.05919879,  0.04225244, ..., -0.23945914,\n",
       "        -0.14030399,  1.31228252]),\n",
       " array([ 0.        , -1.0005614 , -0.89184516, ..., -1.26894611,\n",
       "        -1.00964957, -0.83159191]),\n",
       " array([ 0.        ,  0.70408421,  1.54412846, ..., -3.2284155 ,\n",
       "        -2.7897333 , -2.14465114]),\n",
       " array([ 0.        , -1.58378376, -1.90680166, ..., -1.88280106,\n",
       "        -1.26820058, -0.13192404]),\n",
       " array([ 0.        , -0.68223056, -0.71738807, ..., -2.98058024,\n",
       "        -3.10178986, -3.06540171])]"
      ]
     },
     "execution_count": 19,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "N_norm = rg.randn(num_steps)\n",
    "\n",
    "[OU_process_cy3(num_steps, N_norm=N_norm) for _ in range(num_sim)]\n",
    "[ou_process_cy_mp(num_steps) for _ in range(num_sim)]\n",
    "[ou_process_nb_mp(num_steps) for _ in range(num_sim)]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Parallelize Cython \n",
    "\n",
    "We need to release the GIL. How?\n",
    "\n",
    "- cython functions only callable from C (`cdef`) can release the gil using\n",
    "      cdef function_name() nogi:\n",
    "- cython functions callable from python (`def`, `cpdef`) \n",
    "  can release the gil using a `with nogil:` block.\n",
    "  \n",
    "Since our function needs to be called from python, we use the second approach\n",
    "(see the `OU_process_cy3` definition above)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "metadata": {},
   "outputs": [],
   "source": [
    "num_steps = int(5e6)\n",
    "num_sim = 8"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Parallel cython: OU process"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "285 ms ± 28.8 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)\n"
     ]
    }
   ],
   "source": [
    "%%timeit\n",
    "[OU_process_cy3(num_steps, N_norm=N_norm) for _ in range(num_sim)]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "175 ms ± 42.1 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)\n"
     ]
    }
   ],
   "source": [
    "%%timeit\n",
    "Parallel(n_jobs=4, prefer='threads')(delayed(OU_process_cy3)(num_steps, N_norm=N_norm) \n",
    "                                     for _ in range(num_sim))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Parallel cython: OU process + RNG"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 23,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "793 ms ± 19.1 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)\n"
     ]
    }
   ],
   "source": [
    "%%timeit\n",
    "[ou_process_cy_mp(num_steps) for _ in range(num_sim)]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 24,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "431 ms ± 112 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)\n"
     ]
    }
   ],
   "source": [
    "%%timeit\n",
    "Parallel(n_jobs=4, prefer='threads')(delayed(ou_process_cy_mp)(num_steps) \n",
    "                                     for _ in range(num_sim))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Results\n",
    "\n",
    "- We get a **~2x** speed improvement on a 2-core system \n",
    "- Empirically, best performance with 4 threads on 2-core system"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Parallelize Numba\n",
    "\n",
    "We need to release the GIL. How?\n",
    "\n",
    "- use `numba.jit(nogil=True)`  (see the `OU_process_nb` definition above)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 25,
   "metadata": {},
   "outputs": [],
   "source": [
    "num_steps = int(5e6)\n",
    "num_sim = 8"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Parallel Numba: OU process"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 26,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "269 ms ± 96.2 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)\n"
     ]
    }
   ],
   "source": [
    "%%timeit\n",
    "[OU_process_nb(num_steps, N_norm=N_norm, δt=0.1, x0=0, σ=1, τ=1)\n",
    " for _ in range(num_sim)]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 27,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "159 ms ± 19.3 ms per loop (mean ± std. dev. of 7 runs, 10 loops each)\n"
     ]
    }
   ],
   "source": [
    "%%timeit\n",
    "Parallel(n_jobs=4, prefer='threads')(delayed(OU_process_nb)(num_steps, N_norm=N_norm, δt=0.1, x0=0, σ=1, τ=1)\n",
    "                                     for _ in range(num_sim))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Parallel Numba: OU process + RNG"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 28,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "773 ms ± 40.2 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)\n"
     ]
    }
   ],
   "source": [
    "%%timeit\n",
    "[ou_process_nb_mp(num_steps) for _ in range(num_sim)]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 29,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "403 ms ± 51.7 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)\n"
     ]
    }
   ],
   "source": [
    "%%timeit\n",
    "Parallel(n_jobs=4, prefer='threads')(delayed(ou_process_nb_mp)(num_steps) \n",
    "                                     for _ in range(num_sim))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Results\n",
    "\n",
    "- We get a **~2x** speed improvement on a 2-core system \n",
    "- Empirically, best performance with 4 threads on 2-core system"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Take-home message\n",
    "\n",
    "## Thread-based parallelism\n",
    "\n",
    "*Release the GIL to unlock a performance boost with thread-based parallelism.*\n",
    "\n",
    "Numpy's functions (including `randomgen`) already release the GIL! 🎉🎉🎉\n",
    "\n",
    "## Process-based parallelism\n",
    "\n",
    "Rarely achieves gains unless the application is heavily CPU-bound\n",
    "and has relatively low RAM usage."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Standard lib: Threading and Multiprocessing\n",
    "\n",
    "The `mutiprocessing` package provides a relatively simple API (`Pool`)\n",
    "compared to the `threading` package.\n",
    "\n",
    "Unfortunately, even using `mutiprocessing` is much more cumbersome that joblib."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 30,
   "metadata": {},
   "outputs": [],
   "source": [
    "from multiprocessing import Pool"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 31,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Closing subprocess pool.\n",
      "8.88 s ± 0 ns per loop (mean ± std. dev. of 1 run, 1 loop each)\n"
     ]
    }
   ],
   "source": [
    "%%timeit -n1 -r1\n",
    "with Pool(processes=4) as pool:\n",
    "    try:\n",
    "        Res = []\n",
    "        for i in range(num_sim):\n",
    "            result = pool.apply_async(\n",
    "                ou_process_cy_mp, kwds=dict(num_steps=num_steps, \n",
    "                                            dt=0.1, x0=0, \n",
    "                                            tau=1, sigma=2))\n",
    "            Res.append(result)\n",
    "        X = []\n",
    "        for i, result in enumerate(Res):\n",
    "            X.append(result.get())\n",
    "            #print(f'[Worker {i}] returned an array of shape {len(X[-1])}', flush=True)\n",
    "    except KeyboardInterrupt:\n",
    "        print('\\n>>> Got keyboard interrupt.\\n', flush=True)\n",
    "print('Closing subprocess pool.', flush=True)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The `threading` API is even more complex."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Lessons\n",
    "\n",
    "- Joblib's API rocks!"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3.6 (py36_optimize)",
   "language": "python",
   "name": "py36_optimize"
  },
  "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.6.6"
  },
  "toc": {
   "nav_menu": {},
   "number_sections": false,
   "sideBar": true,
   "skip_h1_title": false,
   "toc_cell": false,
   "toc_position": {
    "height": "591.867px",
    "left": "0px",
    "right": "1003px",
    "top": "111.133px",
    "width": "212px"
   },
   "toc_section_display": "block",
   "toc_window_display": true
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}