Skip to content

fix(io): serialize checkpoints with cloudpickle, falling back to dill - #83

Merged
MRyutaro merged 4 commits into
mainfrom
perf/faster-serializer
Oct 10, 2026
Merged

MRyutaro merged 4 commits into
mainfrom
perf/faster-serializer

Conversation

@MRyutaro

@MRyutaro MRyutaro commented Oct 10, 2026 •

Copy link
Copy Markdown
Owner

背景

チェックポイント保存は dill で書いていますが、dill の Pickler は純 Python 実装なので、小さい Python オブジェクトが大量にあるグループ(list of dict など)で標準 pickle より桁違いに遅くなります。

変更内容

  • elastic_notebook/core/io/serializer.py を追加。保存は cloudpickle を先に試し、失敗したら dill(失敗時は途中まで書いた分を巻き戻してから dill で書き直す)。
  • migrate.py と filesystem_adapter.py の dill.dump を差し替え。読み込みは dill.load のまま(cloudpickle の出力は通常の pickle ストリームなので読める。古い dill 製チェックポイントも読める)。
  • 依存に cloudpickle を追加。

標準 pickle にしなかった理由: ノートブックで定義したクラス・関数(__main__)は pickle だと名前参照で保存されます。保存は成功し、新しいカーネルでの復元で AttributeError になります(下の復元グラフ)。cloudpickle は C 実装の pickler の上で、これらだけ値で保存します。

測定結果

ベンチマークと再現手順は benchmarks/serializer/。Python 3.11 / dill 0.4.0 / cloudpickle 3.1.2、ローカルディスク、bytes へのシリアライズ時間(ファイル書き込みは全候補で同じなので除外)、3 回の中央値。NFS、GPU、torch は未測定。

シリアライズ時間(全 8 候補)

dump time

採用案の dill に対する速度比

speedup

ワークロード dill 採用案 倍率
ndarray 100 MB 0.125 s 0.060 s 2.1x
DataFrame 数値 100 MB 0.124 s 0.060 s 2.1x
DataFrame 文字列 50 万行 0.807 s 0.062 s 13.0x
list of dict (30 万) 4.563 s 0.182 s 25.1x
dict of lists (10 万) 2.182 s 0.041 s 53.0x
ユーザー定義クラスのインスタンス (30 万) 6.214 s 0.574 s 10.8x
lambda (2 万) 0.724 s 0.160 s 4.5x

読み込み時間(変化なし)

load time

別プロセスで復元できるか

restore

試した 8 候補: dill / dill(byref=True) / 標準 pickle / cloudpickle / pickle→例外時 dill / pickle+__main__ 文字列スキャン→dill / 同→cloudpickle / cloudpickle→dill。pickle 系は __main__ のクラス・関数で復元に失敗します。スキャン案は __main__ を含むグループだけ書き直すので、クラスのインスタンスが多いと二重に書いて遅くなります(dill に落とすと 6.5 s)。cloudpickle 単独で十分速く、全ケース復元できたため採用しました。

検証

  • 追加テスト tests/test_serializer.py: 通常データの往復、1 ファイルへの複数連続書き込み、cloudpickle 失敗時のロールバックと dill への切り替え、__main__ のクラスと lambda を別プロセスで復元。
  • uv run pytest: 101 passed, 16 skipped。

未確認

  • NFS 上での速度、pandas 以外(torch / GPU テンソル)の保存速度。
  • dill では保存でき cloudpickle では保存できない型は dill に切り替わる設計だが、網羅はしていない。
  • uv.lock は cloudpickle の追加分だけ手で編集(ローカルの uv が無関係な差分を出すため)。

🤖 Generated with Claude Code

https://claude.ai/code/session_01KeDSzxoCKRTMZX6BxTcYm8

MRyutaro and others added 3 commits October 10, 2026 11:34
dill's Pickler is pure Python and is ~10-50x slower than the C pickler on
groups made of many small Python objects. cloudpickle runs on the C pickler
and still stores __main__ classes/functions by value, so checkpoints restore
in a fresh kernel. Loading is unchanged (dill.load).

Co-Authored-By: Claude Sonnet 5.5 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01KeDSzxoCKRTMZX6BxTcYm8
Co-Authored-By: Claude Sonnet 5.5 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01KeDSzxoCKRTMZX6BxTcYm8
@MRyutaro MRyutaro changed the title perf(io): serialize checkpoints with cloudpickle, falling back to dill fix(io): serialize checkpoints with cloudpickle, falling back to dill Oct 10, 2026
…pter.py

Co-Authored-By: Claude Sonnet 5.5 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01X2LedrYRTf2iaq3R3whSvB
@MRyutaro
MRyutaro marked this pull request as ready for review October 10, 2026 03:32
@MRyutaro
MRyutaro merged commit 3ce4084 into main Oct 10, 2026
7 checks passed
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants