pydanticを大量に使ったらメモリ爆発した件

概要

pydantic.BaseModel を継承したクラスのインスタンスを大量に生成した結果, メモリ上限を突破する問題が発生した. オブジェクト数の推移を確認したところ, set や dict といった一般的なオブジェクトが数十万単位で生成されており, 具体的にどのオブジェクトが大量増加しているかわかりにくい状況だった.

本記事では, pydantic.BaseModel の内部実装に着目し, インスタンスを作るたびにどのような set / dict が保持されるのかをコードレベルで確認する.

その過程を通して, pydantic の使いどころを考えるための材料を共有する.

version 情報

• pydantic: 2.11.3
• python: 3.12.9
• os: windows11

何が起きた?

作っていたもの

数理最適化を用いたフレームワークを開発していた. その際, 決定変数を表すクラスにおいて pydantic.BaseModel を継承していた.
これは, 決定変数を設定する際に下限値が上限値を上回らない, というバリデーションが必要だったため. 単なるデータクラス(@dataclasses.dataclass デコレータの使用)でもよかったが, 「せっかく便利なライブラリ知ってるし使ったろ！」という軽い気持ちで pydantic を使うことにした¹.

実装は大体こんな感じ.

from pydantic import BaseModel, model_validator
...

class DecisionVariable(BaseModel):
    id_: str
    type_: VariableType
    lower_bound: float | None = None
    upper_bound: float | None = None

    @model_validator(mode="after")
    def _validate_lower_ge_upper(self):
        if self.lower_bound is not None and self.upper_bound is not None:
            if self.lower_bound > self.upper_bound:
                error_msg = f"Invalid bounds: lower bound ({self.lower_bound}) > upper bound ({self.upper_bound})"
                raise ValueError(error_msg)
        return self
    ...

起きた事象

上記を実務で使用すると, 本番環境サーバーのメモリが上限に達した².

原因調査

どのオブジェクトが増えているのか, pymplerを使用してオブジェクト数の推移を確認した.

from pympler import muppy, summary

# 処理
...

# どのオブジェクトがどれくらい使っているか取得
all_objects = muppy.get_objects()
sum_obj = summary.summarize(all_objects)
# 文字列に変換して表示
log_text = "\n".join(summary.format_(sum_obj))
print(f"object 総数\n{log_text}")

実行結果は以下³.

types |   # objects |   total size
===== | =========== | ============
  set |      363003 |    154.00 MB
 dict |      388504 |     94.67 MB
 ...

set  count=363003 growth=+113016
dict count=379032 growth=+112798
...

set オブジェクト, dict オブジェクトが数十万単位で大量に作成されている.
そんな一般的なオブジェクトは大量に作成したりしていないのに?

どうやって pydantic が原因だと特定した?

増加した set の中身を objgraph を使用して見たところ, pydantic.BaseModel を継承した DecisionVariableの field 名が大量に列挙された.

import objgraph

...

before_set = set(id(o) for o in objgraph.by_type("set"))

# 処理
...

all_objects = {id(o): o for o in objgraph.by_type("set")}
for key, d in all_objects.items():
    if key not in before_ids:
        print(f"{key}: {d}")

出力は以下.

2476983434624: {'type_', 'upper_bound', 'id_', 'lower_bound'}
2476737019968: {'type_', 'upper_bound', 'id_', 'lower_bound'}
...（以下, 同じような出力が続く）

数理最適化を使ったシステムを実務で作った方ならわかると思うが, 最適化モデリングするにあたり決定変数インスタンスは大量に作られる.
大量に作られた決定変数インスタンス分だけ, field 名の set が作られている状況である.
dict も同様に, 各インスタンスの field 名をキーにその値を格納した辞書となっていた.
メモリを圧迫するのは当然.

何が原因?

pydantic のソースコードを見て確認した.
BaseModel クラスのインスタンスは内部で set と dict を持つようにしているようである.

set

設定されている Attribute の集合を, BaseModel.__pydantic_fields_set__ が持つようにしている.

pydantic/main.py

...
class BaseModel(metaclass=_model_construction.ModelMetaclass):
   ...
   __pydantic_fields_set__: set[str] = _model_construction.NoInitField(init=False)
    """The names of fields explicitly set during instantiation."""
   ...
   @property
    def model_fields_set(self) -> set[str]:
        """Returns the set of fields that have been explicitly set on this model instance.

        Returns:
            A set of strings representing the fields that have been set,
                i.e. that were not filled from defaults.
        """
        return self.__pydantic_fields_set__
   ...

なぜこうなっているのか?

デフォルト値で設定される Attribute は __pydantic_fields_set__ には入らないようである.
なので,

1. クラスに定義された fields
2. インスタンスに設定された fields

を区別するためではなかろうか.

実際, 以下のようなコードを実行すると,

from pydantic import BaseModel


class Entity(BaseModel):
    attr_a: str
    attr_b: int = 10


def main():
    instance_a = Entity(attr_a="a")
    instance_b = Entity(attr_a="b")

    # インスタンスごとに定義される field の集合が, 異なる id を持つ
    set_a = instance_a.model_fields_set
    print(f"set a: {set_a}")
    print(f"id of set a: {id(set_a)}")
    set_b = instance_b.model_fields_set
    print(f"set b: {set_b}")
    print(f"id of set b: {id(set_b)}")

    # attribute を後から設定すると model_fields_set に反映される
    instance_a.attr_b = 100
    set_a_revised = instance_a.model_fields_set
    print(f"set a revised: {set_a_revised}")
    print(f"id of set a revised: {id(set_a_revised)}")


if __name__ == "__main__":
    main()

以下のように出力される.

set a: {'attr_a'}
id of set a: 1620898283584
set b: {'attr_a'}
id of set b: 1620898972160
set a revised: {'attr_a', 'attr_b'}
id of set a revised: 1620898283584

set_a の中身と set_b の中身は同じフィールド名の集合だが, id が異なる. 上述の, pydantic.BaseModel のインスタンスが増えると set オブジェクトが増える, という事象と一致する.

また, 最初の set_a と set_a_revised は中身が異なるが id は同じになっている. 同じオブジェクトとして扱われていることがわかる⁴.

dict

dict オブジェクトが増えたのは, インスタンスごとに __dict__ が定義されているから.
具体的にどこで __dict__ の中身が定義されているかは見つけられなかったが...

pydantic/main.py

...
class BaseModel(metaclass=_model_construction.ModelMetaclass):
   ...
   def __iter__(self) -> TupleGenerator:
        """So `dict(model)` works."""
        yield from [(k, v) for (k, v) in self.__dict__.items() if not k.startswith('_')]
        ...

インスタンスごとに field の値は異なるので, インスタンスごとに辞書が増える理由はわからんでもない.

どう対処した?

DecisionVariable インスタンスの参照を消した⁵結果, 上記の set/dict もメモリから開放された.

教訓

むやみやたらと pydantic でデータクラスを作るべきではない

入り組んだバリデーションロジックが不要なのであれば, @dataclasses.dataclass で十分.
pydantic 使うと余計なオーバーヘッド(メモリに対するもの含む)が発生する上, ライブラリの内部実装まで見に行かなければならない(≒デバッグが苦痛になる)可能性がある.

pydantic の居場所は"境界"がちょうどよいという話もある.

便利なライブラリは「便利なライブラリだから」とむやみに使わず, 置き場所に気を付けよう.

1. これを「オーバーエンジニアリング」という. ↩

2. 開発PCのメモリスペックは十分あったので, 本番にデプロイするまでわからない状況だった. おかげで急ぎの対応が必要になりひどい目にあった. ↩

3. 合計サイズが小さくて「こんなのでメモリ爆発するのかよ」と思われるかもしれない. 実際は, この増加を何度も繰り返すため, 線形的に増加した結果メモリが爆発した, という経緯. ↩

4. 一般に, id関数の値が同じであれば, 同じメモリアドレスを参照し, 同じオブジェクトとして扱われるといえる. こちらも参照. ↩

5. 最初からそうしろよ, と思われそうなので補足する. 作っていたものが, 最適化計算を複数実行し, かつ同じ決定変数を用いていたため, 過去の決定変数オブジェクトの参照を残しておく方が都合が良かった. ↩