AWS LambdaのSnapStart（Python）で初期化処理を試してみる

#AWS Lambda

#Lambda SnapStart

kobayashi.m

2025.01.20

はじめに

データ事業本部のkobayashiです。AWS Lambda SnapStartがPythonで利用可能になりましたが、ランタイムフックを使うことでスナップを取得する前とスナップを復元した後に特定の関数を実行することができます。この挙動を確認してみたのでまとめます。

SnapStartのランタイムフック

ランタイムフックは、Lambdaのスナップショット作成前と再開後に特定のコードを実行できる機能で、主に以下の4つの用途があります：

クリーンアップと初期化
- スナップショット作成前のリソース解放
- 再開後の状態やリソースの再初期化
動的構成
- 環境変更に応じた設定の動的更新
- メタデータの更新
外部統合
- 外部サービスへの通知
- 外部システムの状態更新
- チェックポイントや復元プロセスとの連携
パフォーマンスチューニング
- 依存関係のプリロード
- 起動シーケンスの最適化

Pythonのランタイムフックは、Snapshot Restore for Pythonライブラリで提供される2つのデコレータを使用します。

@register_before_snapshot: スナップショット作成前に実行
@register_after_restore: スナップショット再開時に実行

また、メソッドとしてregister_before_snapshot()とregister_after_restore()も利用可能です。

実行順序は重要で、before_snapshotは登録の逆順、after_restoreは登録順となります。
このパッケージはLambda のランタンタイムにプリインストールされているためデプロイパッケージに含める必要はありません。

Pre-snapshot hookとPost-snapshot hookを試してみる

では実際に以下のコードを使ってその挙動を試してみます。

from snapshot_restore_py import register_before_snapshot, register_after_restore

import uuid

before_snapshot_uuid = ""
after_restore_uuid = ""

before_snapshot_uuid_order = []
after_restore_uuid_order = []

def lambda_handler(event, context):
    # lambda handler code
    lambda_handler_uuid = str(uuid.uuid4())
    uuids = {
        "lambda_handler": lambda_handler_uuid,
        "before_snapshot_uuid": before_snapshot_uuid,
        "after_restore_uuid": after_restore_uuid,
    }
    print(uuids)

    return uuids

def fn_before_snapshot():
    # your logic here
    global before_snapshot_uuid
    before_snapshot_uuid = str(uuid.uuid4())
    before_snapshot_uuid_order.append("fn_before_snapshot")

def fn_after_restore():
    # your logic here
    global after_restore_uuid
    after_restore_uuid = str(uuid.uuid4())
    after_restore_uuid_order.append("fn_after_restore")

register_before_snapshot(fn_before_snapshot)
register_after_restore(fn_after_restore)

やっていることとしては、

スナップショット作成前にfn_before_snapshot変数にuuidを設定
スナップショット再開時にfn_after_restore変数にuuidを設定
lambda_handlerとして実行時にlambda_handler変数にuuidを設定して、3つの変数を表示する
といった処理になります。

ではこのコードでLambdaを作成した後に新しいバージョンを発行しスナップショットを作成します。

$ aws lambda publish-version 
    --function-name PythonLambdaSnapStartInitializationTest

スナップショットを作成には多少時間がかかります。作成完了後に作成したスナップショットを実行してみます。

$ aws lambda invoke \
    --function-name PythonLambdaSnapStartInitializationTest:1 \
    response.txt && cat response.txt
{
    "StatusCode": 200,
    "ExecutedVersion": "5"
}
>{"lambda_handler": "82824406-8cf4-4fb3-97bd-fb8857183b50", "before_snapshot_uuid": "d602f28d-0727-4ebd-b069-9a911281ffa5", "after_restore_uuid": "8eda64f5-e5fa-40ae-a35c-9b6d406e517e"}

続けてもう一度同じスナップショットを実行してみます。

$ aws lambda invoke \
    --function-name PythonLambdaSnapStartInitializationTest:1 \
    response.txt && cat response.txt
{
    "StatusCode": 200,
    "ExecutedVersion": "5"
}
>{"lambda_handler": "7eab43ce-09f7-470b-8122-9431cf6a0e54", "before_snapshot_uuid": "d602f28d-0727-4ebd-b069-9a911281ffa5", "after_restore_uuid": "8eda64f5-e5fa-40ae-a35c-9b6d406e517e"}

実行結果を比べるとlambda_handlerは実行ごとに異なるUUIDとなっています。before_snapshot_uuidとafter_restore_uuidに関しては2回の実行とも同じUUIDとなっています。
before_snapshot_uuidが同一なのはスナップを作成する前にUUIDが設定されたためですが、after_restore_uuidに関しては1回目の直後に2回目の実行を行ったためウォームスタートとなったためです。
しばらく時間をおいて再実行するとbefore_snapshot_uuidだけが同一でlambda_handlerとafter_restore_uuidが前回の実行と異なる値を返します。

$ aws lambda invoke \
    --function-name PythonLambdaSnapStartInitializationTest:1 \
    response.txt && cat response.txt
{
    "StatusCode": 200,
    "ExecutedVersion": "5"
}
>{"lambda_handler": "6af04942-6189-4836-aac4-c63485e5769c", "before_snapshot_uuid": "d602f28d-0727-4ebd-b069-9a911281ffa5", "after_restore_uuid": "d40e1d0e-b035-43aa-9bdd-d8ce2bd61001"}

次に別のスナップショットを作成して実行してみます。

$ aws lambda publish-version 
    --function-name PythonLambdaSnapStartInitializationTest
{
    "FunctionName": "PythonLambdaSnapStartInitializationTest",
    "FunctionArn": "arn:aws:lambda:ap-northeast-1:1234567890:function:PythonLambdaSnapStartInitializationTest:2",
    "Runtime": "python3.12",
$ aws lambda invoke \
    --function-name PythonLambdaSnapStartInitializationTest:2 \
    response.txt && cat response.txt
{
    "StatusCode": 200,
    "ExecutedVersion": "5"
}
>{"lambda_handler": "26b299b1-1acd-4fc3-b98c-4b0f8df0bbea", "before_snapshot_uuid": "902f6b57-d40a-4442-938e-f60e8dff64e0", "after_restore_uuid": "f7a5580c-0bf8-40bf-b836-b37c91c7ae2a"}

新しいスナップショットを実行するとbefore_snapshot_uuidが変わっていることがわかります。

$ aws lambda invoke \
    --function-name PythonLambdaSnapStartInitializationTest:2 \
    response.txt && cat response.txt
{
    "StatusCode": 200,
    "ExecutedVersion": "5"
}
>{"lambda_handler": "d201eadd-83ef-4046-8cc8-6e4597c8fdd0", "before_snapshot_uuid": "902f6b57-d40a-4442-938e-f60e8dff64e0", "after_restore_uuid": "f7a5580c-0bf8-40bf-b836-b37c91c7ae2a"}

再度同じスナップショットを実行すると先程と同じくbefore_snapshot_uuidが変わらないことがわかります。

Pre-snapshot hookとPost-snapshot hookの実行順序

ランタイムフックの実行順序はbefore_snapshotは登録の逆順、after_restoreは登録順となるのでこの挙動を以下のコードを使って試してみます。

from snapshot_restore_py import register_before_snapshot, register_after_restore

before_snapshot_uuid_order = []
after_restore_uuid_order = []

def lambda_handler(event, context):
    # lambda handler code
    orders = {
        "before_snapshot_uuid_order": before_snapshot_uuid_order,
        "after_restore_uuid_order": after_restore_uuid_order,
    }

    return orders

def fn_before_snapshot(order: int):
    # your logic here
    global before_snapshot_uuid_order
    before_snapshot_uuid_order.append(f"fn_before_snapshot_{order}")

def fn_after_restore(order: int):
    # your logic here
    global after_restore_uuid_order
    after_restore_uuid_order.append(f"fn_after_restore_{order}")

register_before_snapshot(fn_before_snapshot, 1)
register_before_snapshot(fn_before_snapshot, 2)
register_after_restore(fn_after_restore, 1)
register_after_restore(fn_after_restore, 2)

ではスナップを作成して実行してみます。

$ aws lambda invoke \
    --function-name PythonLambdaSnapStartInitializationTest:3 \
    response.txt && cat response.txt
{
    "StatusCode": 200,
    "ExecutedVersion": "5"
}
>{"before_snapshot_uuid_order": ["fn_before_snapshot_2", "fn_before_snapshot_1"], "after_restore_uuid_order": ["fn_after_restore_1", "fn_after_restore_2"]}

before_snapshot_uuid_orderを見るとfn_before_snapshot_2,fn_before_snapshot_1とregister_before_snapshot関数で登録した順序と逆になっていることがわかります。
after_restore_uuid_orderに関してはfn_after_restore_1,fn_after_restore_2とregister_after_restore関数で登録した順序になっていることがわかります。

まとめ

AWS LambdaのPython SnapStartでランタイムフックの挙動を確認してみました。
before_snapshotはLambdaがスナップショットを作成する前に実行するコードなので外部ファイルのダウンロードやコンピューティングリソースを大量に使う前処理のような処理で使い、after_restoreはSnapStartがスナップショットが復元されてLambdaハンドラーが実行される前に処理されるコードなのでデータベースとのコネクションや機密情報の取得などで使ったほうが良いです。
また登録順序と実行順序は注意が必要です。この点デコレータを使うよりもフック関数を使ったほうが順序が明確になるので関数の利用をおすすめします。

最後まで読んで頂いてありがとうございました。