Python: NumPy の empty() / zeros() を呼び出した直後は物理メモリの使用量が増えない

表題のとおりなんだけど、NumPy の empty() や zeros() は呼び出した直後はメモリの RSS (Resident Set Size) が増えない。ようするに、呼び出した直後は配列に物理メモリが割り当てられていない、ということ。今回は、そのせいでちょっとハマったのでメモとして詳細について書き残しておく。ただし、この事象が起こるのは仮想記憶に COW (Copy-On-Write) の機構を備えた OS、という条件はある。

使った環境は次のとおり。

$ lsb_release -a
No LSB modules are available.
Distributor ID: Ubuntu
Description:    Ubuntu 20.04.2 LTS
Release:    20.04
Codename:   focal
$ uname -r
5.4.0-1036-gcp

TL; DR

NumPy の empty() / zeros() は内部的に calloc(3) / malloc(3) を呼び出している
- 仮想記憶の COW (Copy-On-Write) が働くので、実際にメモリの書き込みが生じないと物理メモリが割り当てられない
OS がメモリの管理にページング方式を採用していると、使った領域の分だけ物理メモリが増えていく
- そのため zeros() で作ったスパースな配列を操作していると突然物理メモリが爆発して戸惑う
一方で ones() や zeros_like() は malloc(3) した上でメモリに書き込みが生じるため最初から物理メモリが割り当てられる
最初から必要な物理メモリを見積もりたいときは zeros() は使わず empty() した上で 0 を書き込むと良い

下準備

今回は NumPy のソースコードをデバッガで追いかけるので、ちょっと準備の手数が多い。

まずは必要なパッケージをインストールしておく。

$ sudo apt-get -y install git python3-venv gdb

Python の仮想環境を作ってアクティベートする。

$ python3 -m venv example
$ source example/bin/activate

Python の REPL からメモリの情報を取りたいので psutil をインストールしておく。

$ pip3 install psutil

NumPy のリポジトリをクローンして、リリース版のタグをチェックアウトする。

$ git clone https://github.com/numpy/numpy.git
$ cd numpy/
$ git checkout -b v1.20.1 tags/v1.20.1

依存パッケージをインストールしたら、デバッグ用の情報を残して拡張モジュールをビルドする。

$ pip3 install -r test_requirements.txt
$ CFLAGS='-Wall -O0 -g' python3 setup.py build_ext -i

あとは開発モードでインストールする。

$ pip3 install -e .

Python の REPL を起動したら準備完了。

$ python3

実験してみる

それでは、実際に NumPy の empty() や zeros() を呼び出して仮想メモリと物理メモリの使用状況を確認してみよう。

まずは、仮想メモリと物理メモリの使用状況を表示する関数を定義する。

>>> import psutil
>>> def print_memory_usage():
...     process = psutil.Process()
...     mem_info = process.memory_info()
...     print(f'vms={mem_info.vms // 1024 // 1024}MB rss={mem_info.rss // 1024 // 1024}MB')
...

定義したら NumPy をインポートする。

>>> import numpy as np

メモリの初期状態はこんな感じ。

>>> print_memory_usage()
vms=44MB rss=26MB

zeros() を使って配列を作ってみよう。

>>> x = np.zeros((10_000, 10_000))
>>> print_memory_usage()
vms=807MB rss=26MB

すると、仮想メモリの使用量は増えているが、物理メモリの使用量は増えていないことがわかる。

一旦、配列を削除しよう。すると、仮想メモリの使用量も減る。

>>> del x
>>> print_memory_usage()
vms=44MB rss=26MB

empty() も同じように試してみる。

>>> x = np.empty((10_000, 10_000))
>>> print_memory_usage()
vms=807MB rss=26MB

こちらも、同じように仮想メモリの使用量しか増えない。

では、ones() はどうか？

>>> del x
>>> print_memory_usage()
vms=44MB rss=26MB
>>> x = np.ones((10_000, 10_000))
>>> print_memory_usage()
vms=807MB rss=789MB

なんと、これは物理メモリまで増えるのだ。

同じように zeros_like() も、呼び出した直後に物理メモリの割り当てが生じる。

>>> del x
>>> x = np.empty((10_000, 10_000))
>>> print_memory_usage()
vms=807MB rss=26MB
>>> y = np.zeros_like(x)
>>> print_memory_usage()
vms=1570MB rss=789MB

ハマった状況について

この振る舞いによって、実際に自分がハマったパターンについて説明しておく。 zeros() を使って初期化したスパースな配列を操作していたところ、途中までは何も問題がなかった。しかし、配列に対して np.sum() を使ったところ物理メモリが突如として大量に必要となってプロセスが落ちた。

もうお分かりだと思うけど、そもそも初期化した時点で物理メモリに収まらないサイズの配列だった。しかし、スパースだったが故に、途中まで必要とする物理メモリが少なかった。今回であれば np.sum() という、およそメモリを大量に消費することなどなさそうな関数でプロセスが死んでデバッグに苦労した。

問題の回避策

問題を回避するには物理メモリを割り当ててやれば良い。つまり、初期化した後に自分で値を書き込む。例えば、次のように empty() で初期化した上で自分で 0 を書き込んでやる。

>>> del x
>>> del y
>>> print_memory_usage()
vms=44MB rss=26MB
>>> x = np.empty((10_000, 10_000))
>>> print_memory_usage()
vms=807MB rss=26MB
>>> x[...] = 0
>>> print_memory_usage()
vms=807MB rss=789MB

ソースコードを追いかけてみる

ここからは、実際に NumPy のソースコードを追いかけてみよう。

一旦、Python のインタプリタのプロセスは立ち上げ直しておく。

$ python3

そして、改めて NumPy をインポートしよう。

>>> import numpy as np

プロセス ID を確認する。

>>> import os
>>> os.getpid()
14592

別のターミナルで gdb を起動する。

$ sudo gdb

Python の REPL プロセスをアタッチする。

(gdb) attach 14592

NumPy で配列のメモリを確保している箇所にブレークポイントを打つ。

(gdb) break PyDataMem_NEW_ZEROED
Breakpoint 1 at 0x7f4c524c9010: file numpy/core/src/multiarray/alloc.c, line 265.

REPL に制御を戻す。

(gdb) c
Continuing.

REPL で zeros() を呼び出す。

>>> x = np.zeros((100, 100))

ブレークポイントに到達するのでバックトレースを表示する。

Breakpoint 1, PyDataMem_NEW_ZEROED (size=size@entry=80000, elsize=elsize@entry=1) at numpy/core/src/multiarray/alloc.c:265
265    {
(gdb) bt
#0  PyDataMem_NEW_ZEROED (size=size@entry=80000, elsize=elsize@entry=1) at numpy/core/src/multiarray/alloc.c:265
#1  0x00007f4c524c90ba in npy_alloc_cache_zero (sz=80000) at numpy/core/src/multiarray/alloc.c:155
#2  0x00007f4c52500486 in PyArray_NewFromDescr_int (subtype=<optimized out>, descr=0x7f4c527c1940 <DOUBLE_Descr>, nd=nd@entry=2, dims=dims@entry=0x13c1da0, 
    strides=strides@entry=0x0, data=data@entry=0x0, flags=<optimized out>, obj=<optimized out>, base=<optimized out>, zeroed=<optimized out>, allow_emptystring=<optimized out>)
    at numpy/core/src/multiarray/ctors.c:826
#3  0x00007f4c525047ee in PyArray_Zeros (nd=2, dims=0x13c1da0, type=<optimized out>, is_f_order=0) at numpy/core/src/multiarray/ctors.c:2833
#4  0x00007f4c5258dc33 in array_zeros (__NPY_UNUSED_TAGGEDignored=<optimized out>, args=0x7f4c52a894c0, kwds=0x0) at numpy/core/src/multiarray/multiarraymodule.c:2045
#5  0x00000000005f4249 in PyCFunction_Call ()
#6  0x00000000005f46d6 in _PyObject_MakeTpCall ()
#7  0x0000000000570936 in _PyEval_EvalFrameDefault ()
#8  0x000000000056955a in _PyEval_EvalCodeWithName ()
#9  0x000000000068c4a7 in PyEval_EvalCode ()
#10 0x000000000067bc91 in ?? ()
#11 0x000000000067bd0f in ?? ()
#12 0x00000000004a3291 in ?? ()
#13 0x00000000004a4305 in PyRun_InteractiveLoopFlags ()
#14 0x000000000067e059 in PyRun_AnyFileExFlags ()
#15 0x00000000004ee716 in ?? ()
#16 0x00000000006b63bd in Py_BytesMain ()
#17 0x00007f4c530000b3 in __libc_start_main (main=0x4eea30 <main>, argc=1, argv=0x7ffdff0db088, init=<optimized out>, fini=<optimized out>, rtld_fini=<optimized out>, 
    stack_end=0x7ffdff0db078) at ../csu/libc-start.c:308
#18 0x00000000005fa4de in _start ()