python で MemoryError を捕捉する方法

Overview

けっこう昔にこんな投げっぱなしな感じの記事書いたんですけど、
https://qiita.com/arc279/items/d400c92d6019afc665eb

その派生というか、 docker 内のプロセスが OOMKiller に始末されるのを避けて MemoryError を捕捉する方法です。

事前準備と予備知識

とりあえずメモリを500MB に制限して docker で python を起動します。

docker周り

たまたま手元にあったのが python:3.7.6-slim だったので別に深い意味はないです。

docker-compose.yml

version: '3.9'

services:
  node:
    image: python:3.7.6-slim
    volumes:
      - .:/mnt
    deploy:
      resources:
        limits:
          memory: 500M
        reservations:
          memory: 500M
    tty: true
    command: bash

bash

# 起動して
$ docker-compose up

# コンテナに入る
$ docker-compose exec node bash

コンテナ内で確認　

docker のメモリ制限は cgroup で制限されてるので、確認はこう。

bash

root@33fc41fd23a2:/# cat /sys/fs/cgroup/memory/memory.limit_in_bytes
524288000

500MB に制限されてますね。
memory_swap を設定しない場合の swap はメモリ指定の2倍になってるそうです。

ちなみに

free とか /proc/meminfo には反映されないので注意。

ホストから14GB割り当ててる場合の図

root@33fc41fd23a2:/# cat /proc/meminfo  | grep Mem
MemTotal:       14334276 kB
MemFree:        13063124 kB
MemAvailable:   13221092 kB

root@33fc41fd23a2:/# free
              total        used        free      shared  buff/cache   available
Mem:       14334276      468032    12900708      409740      965536    13210688
Swap:       1048572       43808     1004764

python でヒープサイズの上限

以下、 前述のコンテナ内に入っての処理です。
python を起動します。

特に指定しないと無制限になってるので、

python

>>> import resource
>>> resource.getrlimit(resource.RLIMIT_DATA)
(-1, -1)

メモリ 500MB の状況で 4GB くらい確保しようとすると

python

>>> try:
...     _ = bytearray(4000 * 1024 ** 2)
... except Exception as e:
...     print(type(e), e)
...
Killed

except できず問答無用で OOM Killer に始末されて死んじゃいます。

ヒープサイズに制限をかける

で、無言で死んでしまうとちょっと困るので、これをなんとか捕捉したい。

上限を設定します

ヒープサイズの上限を５００MBに設定した！

>>> import resource

>>> resource.setrlimit(resource.RLIMIT_DATA, (524288000, -1))
>>> print(resource.getrlimit(resource.RLIMIT_DATA))
(524288000, -1)

で、再度実行

python

>>> try:
...     _ = bytearray(4000 * 1024 ** 2)
... except Exception as e:
...     print(type(e), e)
...
<class 'MemoryError'>

except で補足できました。

ちなみに MemoryError は無口なので str(e) しても特に文言がありません。
なので type(e) しないと分かりづらいので注意。

上記をまとめると

ヒープサイズの上限を /sys/fs/cgroup/memory/memory.limit_in_bytes の値に設定してやれば良さそうです。

trap_memory_error.py

import resource

# 500MB しかなくても 4GB確保したい
alloc_size = 4000 * 1024 ** 2

with open("/sys/fs/cgroup/memory/memory.limit_in_bytes", "r") as fp:
    memory_limit_in_bytes = int(fp.read())

print("alloc_size", alloc_size)
print("/sys/fs/cgroup/memory/memory.limit_in_bytes", memory_limit_in_bytes)

resource.setrlimit(resource.RLIMIT_DATA, (memory_limit_in_bytes, -1))
print(resource.getrlimit(resource.RLIMIT_DATA))

try:
    _ = bytearray(alloc_size)
    print("OK")
except MemoryError as e:
    print(type(e), e)

bash

root@33fc41fd23a2:/mnt# python trap_memory_error.py
alloc_size 4194304000
/sys/fs/cgroup/memory/memory.limit_in_bytes 524288000
(524288000, -1)
<class 'MemoryError'>

なお、ヒープ上限を設定しないと

上記の resource.setrlimit の行をコメントアウトして実行すると、普通に OOM Killer に始末されて死にます。

root@33fc41fd23a2:/mnt# python no-trap_memory_error.py
alloc_size 4194304000
/sys/fs/cgroup/memory/memory.limit_in_bytes 524288000
(-1, -1)
Killed

っていう話。

プロセスの確保可能なヒープサイズに上限を設定することでプロセス内で捕捉するのが肝なので、
同様の手段があれば python 以外でもできると思います。

cf.

• https://docs.python.org/ja/3/library/resource.html
• https://qiita.com/toshihirock/items/236b7541016929060cc0
• https://yanomy.hatenablog.com/entry/2018/07/31/221957