並列計算したいので,numpyレベルでの並列処理を制限したい.


上のレベル(遺伝的アルゴリズムなどの並列化できるアルゴリズムの実行)で並列計算をするので,下のレベル(行列計算時など)でnumpyが自動的に行っているマルチスレッド処理を止める必要がでてきたので方法を調べましたが結構時間を費やしてしまったので,メモを共有しておきます. Visual Studio Code から実行することを想定しています. (デバッグ機能によるか,jupyterインタラクティブ環境を使うかのどちらか)

LinuxかMac環境で試しており,Windows環境は試していません.

なぜ制限が必要か

コア数32のマシンで並列計算をしたいとします.
一度に30個のプロセスをたててそれらを並列計算にしたいとき,それぞれのプロセスの計算の中でnumpyが自動的に並列計算を行ってしまうと,30プロセスそれぞれが32スレッドの計算をするということが起こり,計30*32=960のスレッドが走ってしまい,コア数の32を大きく超えるため,オーバーヘッドにより計算効率が落ちます.これを防ぐため,下位のnumpyのレベルでのマルチスレッド数を制限しておく必要があります.

わかったこと

numpyといってもどのライブラリを使って並列化しているかは場合により異なる

osの違いやcondaやpipでどこからとってきたか等により,numpyがどの並列ライブラリを使っているかは異なるようです.openblas, mkl, openmp, veclib などの種類があるようです.
どのライブラリを使ってるか調べるには,

print(np.__config__.blas_opt_info['libraries'])

とかで確認できるようです.
参考
https://sgryjp.gitlab.io/posts/2020/2020-09-20/

並列化を制限するには,多くの場合環境変数を設定する必要がある

一部のライブラリの場合,コードの中でもスレッドを制限することができますが,多くの場合環境変数で並列のスレッド数を設定することができるようです.

vscodeからの実行時に環境変数を設定するには,ワークスペース上に.env ファイルを置いておき,そこに環境変数を書き込んでおけば良い.

vscodeからプログラムを動かす場合,プログラムに環境変数を伝えておく必要があります.
.envというファイルをワークスペース(多くの場合ファイルがおいてあるフォルダ)に置いて,中に環境変数の設定を書いておけば,vscodeは実行前にこれを読み込み,制限をかけることができるようになります.vscodeのjupyterのインタラクティブ環境を利用してプログラムを動かす場合や,F5で走るデバッグ機能を利用するときは,これでいけます.ただし,再生ボタンのマークをクリックしてRun Python File in Terminal として動かす場合は,このファイルを読み込んでないようです.

公式 https://code.visualstudio.com/docs/python/environments#_environment-variable-definitions-file

.envファイルの設定例

ひとつの環境で利用する場合,下記すべて設定する必要はないですが,一応書いておきます.

.env
OPENBLAS_NUM_THREADS=1
MKL_NUM_THREADS=1
OMP_NUM_THREADS=1
VECLIB_NUM_THREADS=1
NUMEXPR_NUM_THREADS=1

プログラム

20の並列プロセスを使い,40個の2000x2000行列の固有値を計算する例です.

multi_example.py
import numpy as np
import time
from multiprocessing import Pool

#%%

def eigs(X):
    # 並列処理する関数.固有値を計算する    
    print(f'X.shape={X.shape}')
    ret = np.linalg.eig(X)
    print('done')
    return(ret)

if __name__ == "__main__":
    p = Pool(20)

    matrices = np.random.rand(40,2000,2000)
    time_start = time.time()    
    result = p.map(eigs, matrices)
    print('elapsed time:{}'.format(time.time()-time_start))
    p.close()

実行中に別のターミナルでtopをみて,プロセスのCPU利用率が100%を超えていなかったり,loadがとんでもない値になっていなければ,制限をかけることは成功しています.

ターミナルからの実行

ターミナルからの実行の場合,必要な環境変数のexportを行って

 export OPENBLAS_NUM_THREADS=1

から実行するか,

python multi_example.py

実行時に渡します.

OPENBLAS_NUM_THREADS=1 python multi_example.py