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 parallelization packages available. This wiki will attempt to present minor examples of some common approaches, 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. calling 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, 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 also discuss Python packages that are built 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 previously mentioned, there are a lot of packages for 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/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, one after the next, and asynchronous means multiple tasks can be executed in small parts in overlapping periods. Note that many statistical uses of parallelization do not need synchronous parallelization.

Since both approaches have the same syntax we’ll be using multiprocess (as it works with ipython and jupyter notebooks). multiprocessing can have problems with ipython and jupyter notebooks. In jupyter notebooks, if you want to use multiprocessing , it needs to be wrapped in if name == "main" to properly fork processes; you can also save your process is as a .py file and import it into your jupyter notebook.

Implementation

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

import multiprocess as mp 
#Could also using import multiprocessing as mp; the syntax is the same

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. Just like with multiprocessing, the number of cores needs to be computed before proceeding.

Implementation

import joblib

num_cores = joblib.cpu_count()

To use joblib functionality, we can import the functions Parallel and delayed. Parallel creates an instance for parallelization with a set number of jobs based on the core availability. delayed creates deferred function calls to distribute computations across multiple cores; this is the function that allows for simultaneous task execution.

from joblib import Parallel, delayed

We then can move our list-recursive example into our joblib structure. We wrap Parallel(...)() around 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). This will provide more detailed output information.

   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.

Jupyter Notebook: ipyparallel

For those that heavily use Jupyter notebooks, the package ipyparallel works quite nicely. Some benefits of this package include its verbosity, its ability to serialize numpy arrays cleanly, and the fact that it can distribute across multiple machines. 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 done from the command line with the following command or inside the jupyter notebook extensions tab.

ipcluster nbextension enable

You’ll then 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 them.

1. Passing packages all together

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

2. Passing variables:

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

   dv["my_k"] = my_k

3. We can also use the %px magick to run lines of code on all servers.

   %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

For asynchronous, you need to use .get() in order to get the results at a later time/when they come in, since the work is not done in order and is instead chunked for efficiency.

%%timeit
results = dv.map_async(my_simulation, my_k).get()

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

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

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 |