As my first post — and as a reminder to myself — here’s a quick set of code snippets to use ray, the python library that changed my life when it comes to multiprocessing in Python. Why use ray instead of multiprocessing — I’ll leave that to the plethora of posts online. Here, I provide a simple example on how to use ray— nothing else*!

Getting started

First of all, you’ll need to install ray:

pip install ray

Next, I’ll assume you have numpy installed

A dumb example

Suppose you have the following function that takes the square of a number:

def get_square(x):
    return x**2

If you wanted to square all the values in an array y for the first 100 numbers, you would do:

import numpy as np

n = 100
y = np.arange(n)
result = np.zeros(n)
for i in range(n):

    result[i] = get_square(y[i])

Parallelizing with ray

Let’s now parallelize this same dumb function with ray. To do this, you need to create a new function that will allow you to cast ray’s magic into it — to do this, you simply decorate the function as such:

import ray

@ray.remote
def get_square2(x):
    return x**2

I have called this new magic function get_square2 — but of course, it could be called whatever I wanted. Next, we initialize ray, and define how many cpus we want to use; in this case, I’ll use 4:

ray.init(num_cpus = 4)

And now we run a for loop on which we save executions to our new ray function:

all_results = []
for i in range(n):

    all_results.append( get_square2.remote(y[i]) )

This for loop has not actually ran any calculations. It just staged the calculations to be performed. The next step runs the calculations:

ray_results = ray.get(all_results)

Note this gives back a list with the results — we could pass them back as an array as follows:

result2 = np.zeros(n)
for i in range(n):
    
    result2[i] = ray_results[i]

Et voila!

*: One very interesting note: this example of ray actually produces an example when ray actually performs worse than a normal for loop. ray really excels at functions that are very expensive to evaluate. On simpler ones, like this one, a simple for loop wins.