tempfileライブラリを使った、AWS Lambdaの一時ディスク領域のスマートなクリーンアップ処理

概要

AWS Lambdaのウォームスタートによるバグ防止のための一時ディレクトリのクリーンアップについて、Python標準ライブラリ tempfile を使うとスマートに実装できる。

Lambdaのウォームスタートの仕様と問題

Lambda関数が呼び出されると一時的な実行環境が生成され、一定時間内に同じ関数が呼び出された際は同一の実行環境が再利用される仕様(ウォームスタート)がある。その際、Lambdaの実行環境における書き込み可能領域の /tmp ディレクトリについても、複数の呼び出しで共有される。

各実行環境には、/tmp ディレクトリ内に 512 MB～10,240 MB (1 MB 単位) のディスク領域が用意されています。ディレクトリのコンテンツは、実行環境が停止された際に維持され、複数の呼び出しに使用できる一時的なキャッシュを提供します。

AWS Lambda 実行環境 - AWS Lambda (参照 2022-06-03)

この仕様により、以前の呼び出し時に書き込んだファイルを意図せず読み込んでしまったり、複数回の実行によりディスク領域が不足してしまうなどのバグを生む可能性がある。

os/globライブラリを使ったクリーンアップ処理の実装

上記のような問題を防ぐには、実行終了前に一時ディレクトリ内をクリーンアップする処理を行い、複数の呼び出し間でファイルが共有されないようにする必要がある。
Pythonスクリプトにおいてこの処理を実装すると下記のようになる。

app.py

import os
import glob

def lambda_handler(event, context):

    # 一時ディレクトリ下でファイルの書き込み処理
    tmpfilename = "/tmp/temp.txt"
    with open(tmpfilename, "w") as f:
        f.write("This is temporary file")

    # 一時ディレクトリ配下のファイルを全て削除
    for p in glob.glob("/tmp/*"):
       if os.path.isfile(p):
           os.remove(p)

これで /tmp 配下のファイルは全て削除できるので解決…ではあるがこの処理を毎回書くのはやや冗長にも感じられる。（上記の処理だけではファイルが削除されただけで、配下のディレクトリが残るのも少しモヤモヤする。）

また、ローカルの開発環境がWindowsの場合、ローカルで実行しようとした際にLinuxとWindowsで一時ディレクトリのパスが異なるため、同一コードで実行できないという問題もある。

tempfileライブラリを使ったクリーンアップ処理の実装

標準ライブラリの tempfile の TemporaryDirectory() を使用すると、一時ディレクトリの生成や削除を自動で行ってくれる。

TemporaryDirectory() を使った実装例は下記の通り。

app.py

import tempfile
import os

def lambda_handler(event, context):

    # 一時ディレクトリの生成
    with tempfile.TemporaryDirectory() as tmpdirname:

        # 一時ディレクトリの絶対パスが tmpdirname に代入され、ディレクトリの存在も確認できる。
        print(tmpdirname) # '/tmp/tmp5c4nrfhm'
        print(os.path.isdir(tmpdirname)) # True

        # 一時ディレクトリ下でファイルの書き込み処理
        tmpfilename = os.path.join(tmpdirname, "temp.txt")
        with open(tmpfilename, "w") as f:
            f.write("This is temporary file")

    # with句を抜けると、生成された一時ディレクトリは削除される。
    print(os.path.isdir(tmpdirname)) # False

    # 当然、一時ディレクトリ内に保存したファイルも削除される。
    print(os.path.isfile(tmpfilename)) # False

Lambdaの実行環境の場合、一時ディレクトリは /tmp/tmp******** (*は乱数)に生成され、このパスが TemporaryDirectory() の戻り値になる。このディレクトリパスは実行環境に合わせて自動で設定されるので、Windowsで実行する際も同一コードで対応可能である。

実行中の環境においてどのディレクトリに一時ディレクトリが作成されるかは、gettempdir() メソッドで確認できる。

print(tempfile.gettempdir()) # '/tmp'

一時ディレクトリが生成される場所については、ドキュメントにもルールの記載がある。

1. 環境変数 TMPDIR で与えられているディレクトリ名。
2. 環境変数 TEMP で与えられているディレクトリ名。
3. 環境変数 TMP で与えられているディレクトリ名。
4. プラットフォーム依存の場所:
   1. Windows ではディレクトリ C:\TEMP 、 C:\TMP 、 \TEMP 、および \TMP の順。
   2. その他の全てのプラットフォームでは、 /tmp 、 /var/tmp 、および /usr/tmp の順。
5. 最後の手段として、現在の作業ディレクトリ。

tempfile --- 一時ファイルやディレクトリの作成 — Python 3.10.4 ドキュメント (参照 2022-06-03)

まとめ

tempfile を利用することでコードの記述量が減り、Windows等の異なる環境でも実行できるソフトコーディングな実装ができる。
また、一時ディレクトリが必要となるスコープも明確になるため、可読性の向上も期待できる。