Parallelization in python

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):
    """
    estimate 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))

import time 

starttime = time.time()
[my_simulation(k) for k in my_k]
print(str(time.time() - starttime) + " seconds")

0.9454693794250488 seconds

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

starttime = time.time()
with mp.Pool(processes=num_cores-2) as p:
  results = p.map(my_simulation, my_k)
print(str(time.time() - starttime) + " seconds")

0.7638578414916992 seconds

Asynchronous

starttime = time.time()
with mp.Pool(processes=num_cores-2) as p:
    results = p.map_async(my_simulation, my_k).get() 
print(str(time.time() - starttime) + " seconds")

0.8183681964874268 seconds

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.

starttime = time.time()
r = Parallel(n_jobs=num_cores-2)(delayed(my_simulation)(k) for k in my_k)
print(str(time.time() - starttime) + " seconds")

0.5306367874145508 seconds

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.

ray

The package ray is an open source project for parallel and distributed computing in Python. ray is often preferred over multiprocess/multiprocessing because of its efficient machine error handling methodology and its ability to handle larger data types. ray allows for aggregation of results/values, meaning the output of one task can be used as input for subsequent tasks seamlessly. The implementation for ray is not dissimilar from that of multiprocess.

Implementation

import ray
import os

In order to parallelize computation with ray, you must first initialize an instance of ray. Just like with multiprocess, the instance relies on the number of CPUs available.

ray.init(num_cpus = os.cpu_count())

Additionally, in order for ray to work, the function my_simulation needs to have the decorator @ray.remote prior to its heading.

@ray.remote
def my_simulation(k):
    """the rest of the function code goes here"""

Once we have initialized an instance based on our number of available CPUs, and we have made sure our function has the proper decorator, we can parallelize computation as follows. The .remote() passes the variable into the function, even if the previous computation has not finished running, and the .get() compiles the results once finished.

starttime = time.time()
futures = [my_simulation.remote(k) for k in my_k]
results = ray.get(futures)
print(str(time.time() - starttime) + " seconds")

This .remote() piece is useful because you can pass previous computations into it in order to have a parallelized, autoregressive work flow.

0.13494515419006348 seconds

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

Just-in-time compilation is the compiling of your computer code during execution instead of beforehand. This means bytecode is not translated to machine code until run time, and is then executed directly. Code is continuously analyzed to determine where compilation/recompilation would speed up performance. Used in conjunction with parallelization, this is an efficient way to improve computational speed.

numba

numba is a high performance Python compiler that translates Python/numpy code into fast, efficient machine code. Python is traditionally regarded as a "slow" language; using numba, you can approach the speeds of languages such as C. numba is specifically designed to work well with numpy arrays.

Implementation

Similar to ray, numba requires a decorator @njit (parallel = True) before the function. We will define a different function to be parallelized here to demonstrate numba prange which allows for efficient parallelization across loops.

A = np.random.rand(3, 3)
B = np.random.rand(3, 3)

@njit(parallel = True)
def my_simulation(A, B):
    C = np.empty_like(A)
    for i in prange(A.shape[0]):
        for j in range(A.shape[1]):
            C[i, j] = A[i, j] + B[i, j]
    return C

prange should not be nested and should only ever be used in the outer loop. This is a very simple example, but for computationally complex looping problems, numba is highly efficient. However, you need to ensure that you are not creating or declaring a variable within a loop that cannot be accessed outside of the loop, or else parallelization will not work.

Additional Topics and Resources

1. Duke’s statistical computing course provides a set of examples for multiple Python parallelization packages