Parallelization in python

This article is still being developed - released earlier to align with Department seminar.

In python, parallelization is a bit more of a “wild west” in terms of the number of packages that can help you parallelize your processes (This wiki will attempt to present minor examples of some along with an outline of things to think about in selecting the library you end up using. In general, the process in python involves

1. defining a function that captures what you want to have parallelized
2. call that function repeatedly in the the parallelization package way (which often looks like a list-recursive call).

There are a few different general frameworks to accomplish this.

1. general wrappers, define the cores you’ll be using and then wrap that around the function call in a list-recursive way
2. just in time compilers, in addition to providing parallelization tools these packages provide decorators (e.g. @jit) for functions that allow for just in time compilation of the python function into lower level machine code (like C++)

We will discuss a python package that is build specifically to parallelize jobs inside jupyter notebooks.

Non-parallelized: single core example

We naturally want to compare a parallelization approach with the base single core analysis, which can be done with the code as follows:

#base for-loop
import numpy as np
from sklearn.linear_model import LinearRegression

my_k = list(np.arange(14)+1)*20

def my_simulation(k):
    """
    estamine the polynomial fit of degree k's MSE for a linear
    relationship

    Parameters
    ----------
    k : int
    degree of polynomial fit

    Returns
    -------
    mse of fit based on 30 points applied to a test set of 10000 points
    """
    x = np.random.uniform(size=30,low=-3,high=3)
    y = 2*x + 2* np.random.normal(size=30)
    y = y.reshape(-1,1)
    X_poly = np.concatenate([x.reshape(-1,1)**kk
                          for kk in np.arange(1,k+1,dtype = int)],
                          axis = 1)


    x0 = np.random.uniform(size=10000,low=-3,high=3)
    y0 = 2*x0 + 2*np.random.normal(size=10000)
    y0 = y0.reshape(-1,1)
    X0_poly = np.concatenate([x0.reshape(-1,1)**kk
                          for kk in np.arange(1,k+1,dtype = int)],
                          axis = 1)

    my_lm = LinearRegression().fit(X_poly,y)
    y0_hat = my_lm.predict(X0_poly)

    return(np.mean((y0-y0_hat)**2))

%%timeit -n3

[my_simulation(k) for k in my_k]

676 ms ± 37.6 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)

General python parallelization packages

As mentioned there are a lot of package to do general parallelization. Two commonly used ones are joblib and multiprocessing. joblib looks a bit more like list-recursion and multiprocessing uses functions like .map(). multiprocessing can have problems with ipython and jupyter notebooks (though joblib wraps it and multiprocess is just a different fork of multiprocessing).

multiprocessing/multiprocess

The multiprocessing package is actually built into python, and works in a pretty straightforward manner. It also allows for synchronous vs asynchronous parallelization, where synchronous means the order of the jobs must be followed and asynchronous where the order of jobs need not be followed. Many statistical uses of parallelization does not need synchronous parallelization.

There since both approaches have the same syntax we’ll be using multiprocess (as it works with ipython and jupyter notebooks).

parallelization

First, you need to identify the number of cores you have:

import multiprocess as mp
#import multiprocessing as mp

num_cores = mp.cpu_count()

synchronous

%%timeit

with mp.Pool(processes=num_cores-2) as p:
  results = p.map(my_simulation, my_k)

209 ms ± 24.7 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)

asynchronous

%%timeit

with mp.Pool(processes=num_cores-2) as p:
    results = p.map_async(my_simulation, my_k).get() # not clear why I needed to use .get() to get something returnable for async

48.9 ms ± 3.08 ms per loop (mean ± std. dev. of 7 runs, 10 loops each)

joblib

The joblib package is pretty straightforward and requires you to structure things in terms of list-recursion with a few wrappers to help the process work as desired.

setup

import joblib

num_cores = joblib.cpu_count()

implimentation

To use joblib functionality we can import the functions Parallel and delayed.

from joblib import Parallel, delayed

We then can move our list-recursive example into our joblib structure. We wrap Parallel(...)() around the old of the inside of the list-recursion and delayed() around the functions that will be repeated multiple times.

%% timeit

r = Parallel(n_jobs=num_cores-2)(delayed(my_simulation)(k) for k in my_k)

172 ms ± 23.4 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)

Add-ons

There are multiple add-ons that can be desirable

1. verbosity can be added with a single parameter use (as seen below) - the large the value, the more messages:

   r = Parallel(n_jobs=num_cores-2, verbose=10)(delayed(my_simulation)(k) for k in my_k)

2. multiple parameters can be passed with redefinitions of functions to take groupings. To replace multiple for-loops often the function itertools.product() is used for the parameters.

Just in Time compilers + parallelization

TODO: numba and cython. The idea of these is to compile the function you will be using the first time you use it and then never compile it again (speeding up the actual use of the function). That is then combined with standard parallelization

Jupyter notebook: ipyparallel

For those that heavily use jupyter notebooks and nothing else (which this author personally doesn’t encourage), the package ipyparallel nicely works inside jupyter notebooks (but so does many of the other packages above like joblib, multiprocess, etc.). One nice this about this package is the verbosity, etc. works intellegently in the notebook. This package also works with numba.jit and cython.

Before opening your jupyter notebook you’ll want to make sure the extension is enabled - which can be down from the command line with the following command or inside the jupyter notebook extensions tab.

ipcluster nbextension enable

And then you’ll need to go to the “IPython Clusters” tab in jupyter noteook and start some parallel clusters.

ipyparallel parallelization

import ipyparallel as ipp

For direct execution we define dv and provide it with what it needs…

rc = ipp.Client()
dv = rc[:] # direct view

passing things to each cluster

ipyparallel provides ways for the user to pass only the necessary items to the the clusters (as to not overload the clusters with too many things).

1. passing packages all together

   with dv.sync_imports():
       from sklearn.linear_model import LinearRegression
       # getting dependencies on the clusters

2. passing variable over:

   Below is just an example - we don’t need to pass my_k since dv will directly get it when we wnat to use it.

   dv["my_k"] = my_k

3. we can also use the %px magick to run lines of code on all servers. Given I think numpy is cooler as np I actually have to do it this way, not the dv.sync_imports way.

   %px import numpy as np

actual functionality

The actual functionality is very similar to multiprocess/multiprocessing with maps and apply.

synchronous

%%timeit
results = dv.map_sync(my_simulation, my_k)

201 ms ± 34.1 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)

asynchronous

%%timeit
results = dv.map_async(my_simulation, my_k).get() # not clear why I needed to use .get() to get something returnable for async

209 ms ± 28.1 ms per loop (mean ± std. dev. of 7 runs, 10 loops each)

additional topics and resources

1. Ray python package allows paralleilization across multiple computers
2. Duke’s statistical computing course provides a set of examples for multiple python parallelization packages |