Repository navigation
fix(io): serialize checkpoints with cloudpickle, falling back to dill - #83
Merged
Merged
Conversation
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
…pter.py Co-Authored-By: Claude Sonnet 5.5 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01X2LedrYRTf2iaq3R3whSvB
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
背景
チェックポイント保存は 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 候補)
採用案の dill に対する速度比
読み込み時間(変化なし)
別プロセスで復元できるか
試した 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。未確認
uv.lockは cloudpickle の追加分だけ手で編集(ローカルの uv が無関係な差分を出すため)。🤖 Generated with Claude Code
https://claude.ai/code/session_01KeDSzxoCKRTMZX6BxTcYm8