Setting CPU Affinity with IPython.parallel and psutil

CPU Affinity lets you pin processes to particular processors on your machine. The Python package psutil lets you view and edit the CPU affinity of any process.

In [4]:
import os
import psutil
p = psutil.Process(os.getpid())
p.get_cpu_affinity()
Out[4]:
[0, 1, 2, 3]

Setting CPU affinity can be a problem if you aren't aware that it is happening. The most common way for CPU affinity to be set without the user being aware is importing numpy linked against particular BLAS libraries.

The most common symptom of this will be that all of your IPython engines will use just one core.

Let's check if this happens here:

In [5]:
import numpy
p.get_cpu_affinity()
Out[5]:
[0, 1, 2, 3]

Okay, so we are fine. There are a few logical cases where we want to set the CPU affinity of our engines.

  1. We want to pin each engine to a particular CPU core.
  2. Something else set the affinity to something we don't want, and we want to put it back to the default.
  3. We want to limit all of our engines to a few CPUs

Pinning each engine to a core

In [9]:
from IPython import parallel
In [11]:
rc = parallel.Client()
dview = rc[:]
rc.ids
Out[11]:
[0, 1, 2, 3]

This assumes that you already only have one engine per core. In this case, I have four.

In [14]:
dview.scatter("cpu", range(len(dview)), flatten=True)
%px print cpu
[stdout:0] 0
[stdout:1] 1
[stdout:2] 2
[stdout:3] 3

Now we can set the affinity of each engine to a single CPU core, so the engines won't be fighting over CPU resources.

In [18]:
%%px
import os
import psutil

p = psutil.Process(os.getpid())
p.set_cpu_affinity([cpu])
print p.get_cpu_affinity()
[stdout:0] [0]
[stdout:1] [1]
[stdout:2] [2]
[stdout:3] [3]

Unpinning engines

A simpler case is just restoring CPU affinity after some library may have set it inappropriately.

We use the multiprocessing.cpu_count to get the number of CPUs.

In [23]:
%%px
import os
import psutil
from multiprocessing import cpu_count

p = psutil.Process(os.getpid())
p.set_cpu_affinity(range(cpu_count()))
print p.get_cpu_affinity()
[stdout:0] [0, 1, 2, 3]
[stdout:1] [0, 1, 2, 3]
[stdout:2] [0, 1, 2, 3]
[stdout:3] [0, 1, 2, 3]