Pythonのpickleファイルが壊れないようにatomicに処理する

https://github.com/shibayu36/chat-hatenablog を使っていて、途中でC-cなどでプロセスをterminateした時にpickleファイルが壊れることがあった。pickleファイルが壊れると、次のpickle.load時に EOFError: Ran out of input のようなエラーが出て、再度作り直しとなってしまう。これはかなり辛い。

色々調べているとファイル操作をatomicになるようにすることで、pickleファイルを壊れないようにできたので、そのやり方をメモしておく。間違えている部分があれば指摘してください。

なぜ壊れるか

たとえば pickle.dump(data, f) というコードがあるとすると、pickleはファイルにdataをserializeしたものを書き込む。この操作はファイルに逐次書き込みをするため、この処理が行われている間にTERM signalが送られてプロセスが死んだ時には、中途半端な状態となってしまう。

その後にそのファイルからpickle.loadしたら、中途半端なところまで読み込みをした後、完全なデータロードが終わる前にEOFに到達する。このためEOFError: Ran out of inputとなる。

どうすると良いか

壊れたファイルが作成されるのを防ぎたい。具体的には以下のようにすると良いようだ。

with tempfile.NamedTemporaryFile(mode="wb", delete=False) as f:
    # pickle.dumpが途中で止まると壊れるのは防げないので、tempileに書き込む
    pickle.dump(data, f)
    # (1) f.flush -> os.fsyncを使い、データが確実にファイルに書き込まれた状態にする
    f.flush()
    os.fsync(f.fileno())
    temp_file_path = f.name

# (2) os.replaceでatomicにファイル操作
os.replace(temp_file_path, file_path)

基本的にはpickle.dumpが途中で止まって壊れるのは防げないので、tempfileに書き込んだ後os.replaceでatomicにファイル移動をするという作戦となる。カッコ書きをした部分をもう少し詳細に説明する。

(1)の部分はファイルに確実に書き込まれた状態にするためのコードだ。確実に書き込まれた状態にするには、Python側で持つバッファとOS側で持つバッファの両方をファイルに書き出さなければならない。f.flushはPython側、os.fsyncはOS側の書き出しだ。次のURLが参考になる。

(2)の部分は、完全に書き込まれたpickleファイルを意図した場所にatomicに移動するコードだ。Pythonの場合atomicに移動する方法としてos.renameとos.replaceがあるようだが、Python 3.3以上なら基本的にos.replaceを使えば良さそうだ。

os.renameは宛先にファイルがもともと存在していたときの挙動が、Windowsの時とそれ以外の時で異なる
os.replaceはクロスプラットフォームで扱える。ただしPython 3.3で追加されたので、それ以前のバージョンをサポートしたい場合は利用できない

ちなみにどんな場合でも確実にtempfileを消したいと思った場合、これら全体をtryで囲み、tempfileを消す必要がある。今回はそこまで踏み込まない。

これら調査には次のリンクを参考にした。

参考: 今回の調査はChatGPT(GPT-4)と一緒に行った

「pickleファイルが壊れないように」という目的は決まっていたので、その目的に向けてChatGPTを使った。このときの調べかたも参考になるかと思い、参考としてつけておく。

まずシンプルに、pickle.dumpで壊れるんだけどどうしたらいい？という質問から始めた。

しかしこのコードが正しいのかの判定が自分にはできなかった。そこで、ChatGPTへそのままこのコードを渡し、問題点はどこ？と聞いてみた。するとflushとfsyncをしないといけないと教えてもらえた。そのままコードを渡したら、より良いコードが出てくるというのはすごい。

↓

ここまで来ると、調査のためのとっかかりが掴めている。つまり、flushやfsyncを呼び出すことや、os.replaceを呼び出すことがなぜ必要かについて自分でも調査し、ダブルチェックをすれば良い。Google検索やGitHubのコード検索を用いて、公式ドキュメントでの記載があることやGitHub上で普通に行われていることを確認する。

まずGoogle検索では以下のものが見つかり、公式ドキュメントでもfsyncが必要なことが言及されていることを確認。

GitHubで検索をかけると、使われていることは確認できたが、os.renameの方が多く使われていることがわかった。

何でos.replaceを提案したのか気になったので、これもChatGPTへ聞いてみる。

↓

os.replace and os.rename are both used to rename files or directories in Python. However, there is a significant difference between the two:

os.rename(src, dst): This function renames the file or directory from src to dst. If dst already exists, it will raise an OSError on Windows and silently overwrite the destination on Unix-based systems. This behavior can lead to data loss if you're not careful.

os.replace(src, dst): This function also renames the file or directory from src to dst. However, if dst already exists, it will be replaced regardless of the operating system. This function was introduced in Python 3.3 to provide a consistent behavior across different platforms.

In summary, os.replace provides a more consistent behavior across different operating systems when it comes to renaming files and directories, especially when the destination already exists. It is generally recommended to use os.replace over os.rename for better cross-platform compatibility.

なるほどと思い、実際に公式ドキュメントを見てみると、https://docs.python.org/ja/3/library/os.html#os.rename あたりにreplaceのcross-platform性に関する記載がされていた。またPython 3.3以降に導入されたことを知り、GitHub上でrenameが多い理由にも納得ができた。

このような調査をして、大体合っていそうだという気持ちになった。