トップ > AWS re:Invent 2024 > PythonでもAWS Lambda SnapStartでコールドスタートを改善することができるようになりました

PythonでもAWS Lambda SnapStartでコールドスタートを改善することができるようになりました

AWS re:Invent 2024 AWS Lambda Python
記事タイトルとURLをコピーする

こんにちは。
アプリケーションサービス部、DevOps担当の兼安です。

今回は、先日AWS Lambdaのコールドスタートを短縮するAWS Lambda SnapStartがPythonのサポートを開始したということで、これについてご紹介します。

本記事の対象者

本記事はAWS Lambdaの基礎知識がある方、そして需要増に備えてチューニングの必要性を感じておられる方を対象としています。

AWS Lambda SnapStartとは

AWS Lambda SnapStartは、AWS Lambda関数のコールドスタートの時間を短縮する機能です。
これまでは、Javaのみで利用可能でしたが、Pythonと.NETでも利用可能になりました。

aws.amazon.com

AWS Lambda SnapStartについては、AWS re:Invent 2024のセッションでも取り上げられており、動画も公開されています。
併せて視聴いただけると、より理解が深まるかと思います。

youtu.be

AWS Lambdaのコールドスタートとウォームスタート

AWS Lambda SnapStartの紹介の前に、まずはAWS Lambdaのコールドスタートとウォームスタートについて簡単に説明します。
コールドスタートとは、Lambda関数が初めて呼び出された際に、Lambdaランタイムが新しいコンテナを起動し、関数のコードを読み込んで初期化することを指します。
この初期化処理には、Lambda関数のコードを読み込むだけでなく、依存関係の解決や初期化処理などが含まれるため、コールドスタートの時間はウォームスタートに比べて長くなります。
反対に、ウォームスタートとは、Lambda関数が既存のコンテナ上で素早く実行される状態を指します。

aws.amazon.com

一般的に、コールドスタートは以下のようなケースで発生します。

  • AWS Lambda関数のコードや設定が更新された場合
  • 需要増により、AWS Lambda関数の新しい実行環境が必要になった場合

基本的に、AWS Lambda関数を一度呼び出した後は、極力再利用がされるようになっており、ウォームスタートが発生することが多いですが、需要増によりコールドスタートが発生することもあると解釈した方が良いでしょう。

AWS Lambdaのコールドスタートを起こしてみる

まずは、Lambda関数におけるコールドスタート自体を確認しておきます。
サンプルコードはこちらです。

import time
import logging

logger = logging.getLogger()
logger.setLevel(logging.INFO)

large_resource = None

def initialize_large_resource():
    global large_resource
    logger.info("Initializing large resource (only on cold start)...")
    time.sleep(3)  # 擬似的な遅延
    large_resource = [i for i in range(10**6)]  # 大きなリストを作成
    logger.info("Resource initialized!")

# Lambdaハンドラー
def lambda_handler(event, context):
    global large_resource

    # グローバル変数が初期化されていないなら初期化
    if large_resource is None:
        initialize_large_resource()

    start_time = time.time()
    response = {"message": "Hello, World!", "input": event}

    end_time = time.time()
    logger.info(f"Execution time: {end_time - start_time:.3f} seconds")

    return response

このコードをデプロイした直後に、Lambda関数を呼び出すと、コールドスタートが発生します。
初回はグローバル変数が初期化されていないので、初期化処理が実行されます。
初期化処理に数秒かかる実装にしているため、起動時に時間がかかる挙動になります。

Function Logs:
START RequestId: 6bbce71b-04cd-4f2f-b482-bedd5693eb0e Version: $LATEST
[INFO]  2024-12-24T09:05:00.610Z    6bbce71b-04cd-4f2f-b482-bedd5693eb0e    Initializing large resource (only on cold start)...
[INFO]  2024-12-24T09:05:04.812Z    6bbce71b-04cd-4f2f-b482-bedd5693eb0e    Resource initialized!
[INFO]  2024-12-24T09:05:04.812Z    6bbce71b-04cd-4f2f-b482-bedd5693eb0e    Execution time: 0.000 seconds
END RequestId: 6bbce71b-04cd-4f2f-b482-bedd5693eb0e
REPORT RequestId: 6bbce71b-04cd-4f2f-b482-bedd5693eb0e  Duration: 4241.56 ms    Billed Duration: 4242 ms    Memory Size: 128 MB Max Memory Used: 69 MB  Init Duration: 101.13 ms

Request ID: 6bbce71b-04cd-4f2f-b482-bedd5693eb0e

2回目以降の実行では、ウォームスタートが発生し、初期化処理がスキップされるため、実行時間が短縮されます。

Function Logs:
START RequestId: 874e1347-0313-4a38-b163-050a02eb87c4 Version: $LATEST
[INFO]  2024-12-24T09:05:48.614Z    874e1347-0313-4a38-b163-050a02eb87c4    Execution time: 0.000 seconds
END RequestId: 874e1347-0313-4a38-b163-050a02eb87c4
REPORT RequestId: 874e1347-0313-4a38-b163-050a02eb87c4  Duration: 1.81 ms   Billed Duration: 2 ms   Memory Size: 128 MB Max Memory Used: 69 MB

Request ID: 874e1347-0313-4a38-b163-050a02eb87c4

この関数に対して、改行を追加でもなんでもよいので修正を加えてデプロイし、再度呼び出すと、再びコールドスタートが発生します。
このように、コールドスタートはLambda関数のコードや設定が更新された場合に発生することがわかります。

AWS Lambda SnapStartを使ってコールドスタートを改善

では、これに対してAWS Lambda SnapStartを試してみます。
SnapStartは以下のステップで動作することで、コールドスタートの時間を短縮する機能です。

  1. Lambda関数のコードをスナップショットとして保存
  2. 次回以降の起動時にスナップショットを復元する
  3. 復元後の状態からLambda関数を起動する

この時、1.と2.において処理を挟むことができ、これをフックと呼びます。
このフックと併せることで、より実務に沿った使い方がイメージできると思うので、一気にフック処理を試してみます。
PythonにおけるSnapStartのフックの使い方はこちらに記載されています。

docs.aws.amazon.com

フック：register_before_snapshotを使ってみる

まずは、SnapStartのフック：register_before_snapshotを使ってみます。

import time
import logging
from snapshot_restore_py import register_before_snapshot

logger = logging.getLogger()
logger.setLevel(logging.INFO)

large_resource = None

@register_before_snapshot
def before_checkpoint():
    if large_resource is None:
        initialize_large_resource()

def initialize_large_resource():
    global large_resource
    logger.info("Initializing large resource (only on cold start)...")
    time.sleep(3)  # 擬似的な遅延
    large_resource = [i for i in range(10**6)]  # 大きなリストを作成
    logger.info("Resource initialized!")

# Lambdaハンドラー
def lambda_handler(event, context):
    global large_resource

    # グローバル変数が初期化されていないなら初期化
    if large_resource is None:
        initialize_large_resource()

    start_time = time.time()
    response = {"message": "Hello, World!", "input": event}

    end_time = time.time()
    logger.info(f"Execution time: {end_time - start_time:.3f} seconds")

    return response

これをデプロイした後、スナップショットを作成します。
まずは、設定でSnapStartを有効にします。

SnapStartの有効化

この状態でバージョンを作成します。
メッセージもあるように、SnapStartを有効にしてからバージョンを作成すると、バージョン作成に時間がかかるようになります。

ナップショットを有効にしてからバージョンを作成すると、バージョン作成に時間がかかるようになる

晴れてバージョンが作成されたら、そのバージョンを指定してLambda関数を呼び出すと、初回のコールドスタート時でもグローバル変数の初期化がスキップされ、起動時間が短縮されていることが確認できます。
これはSnapStartの機能により、コードの読み込みとグローバル変数の初期化処理が済んだ状態でLambda関数が起動されるためです。
フック：register_before_snapshotは、スナップショットを作成する前に実行されるため、初期化処理がスナップショットに含まれることになります。

Function Logs:
RESTORE_START Runtime Version: python:3.13.v13  Runtime Version ARN: arn:aws:lambda:ap-northeast-1::runtime:b881cbc9a10a8bcb3def9d9e9fe38f922bb36510a1d92d4ce85cf2a899eeabd8
RESTORE_REPORT Restore Duration: 566.08 ms
START RequestId: 23361531-c5e9-4080-9978-0d9be9793e03 Version: 1
[INFO]  2024-12-24T08:51:45.369Z    23361531-c5e9-4080-9978-0d9be9793e03    Execution time: 0.000 seconds
END RequestId: 23361531-c5e9-4080-9978-0d9be9793e03
REPORT RequestId: 23361531-c5e9-4080-9978-0d9be9793e03  Duration: 28.67 ms  Billed Duration: 93 ms  Memory Size: 128 MB Max Memory Used: 76 MB  Restore Duration: 566.08 ms Billed Restore Duration: 64 ms

Request ID: 23361531-c5e9-4080-9978-0d9be9793e03

フック：register_after_restoreを使ってみる

今度は、スナップショットを復元した後のフック：register_after_restoreを使ってみます。

import time
import logging
from snapshot_restore_py import register_after_restore

logger = logging.getLogger()
logger.setLevel(logging.INFO)

large_resource = None

@register_after_restore
def after_restore():
    if large_resource is None:
        initialize_large_resource()

def initialize_large_resource():
    global large_resource
    logger.info("Initializing large resource (only on cold start)...")
    time.sleep(3)  # 擬似的な遅延
    large_resource = [i for i in range(10**6)]  # 大きなリストを作成
    logger.info("Resource initialized!")

# Lambdaハンドラー
def lambda_handler(event, context):
    global large_resource

    # グローバル変数が初期化されていないなら初期化
    if large_resource is None:
        initialize_large_resource()

    start_time = time.time()
    response = {"message": "Hello, World!", "input": event}

    end_time = time.time()
    logger.info(f"Execution time: {end_time - start_time:.3f} seconds")

    return response

こちらも同じように、これをデプロイした後、スナップショットを有効化した状態でバージョンを作成します。
そして、作成したバージョンを指定してLambda関数を呼び出すと、初回はグローバル変数の初期化処理が走り起動が遅く、2回目以降は初期化処理がスキップされて起動が早いというスナップショットを適用する前と同じ挙動を確認できます。
フック：register_after_restoreは、スナップショットを復元した後に実行されるため、初期化処理がスナップショットに含まれないことになります。
これが、register_before_snapshotを使った時との挙動の違いを産んでいます。

register_before_snapshotとregister_after_restoreの使い分け

Lambda関数を書く時、起動を早くするため一部の変数をグローバル変数にしておくことがあります。
この時、データベースコネクションなど、永続的に保持されると問題になる変数は、register_after_restoreを使うとよいでしょう。
それら以外の、1回しか初期化しない変数・データは、register_before_snapshotを使うとよいと思われます。

まとめ

AWS Lambda SnapStartを使うことで、Lambda関数のコールドスタートを改善することができます。
コールドスタートはデプロイ以外だと、不意な需要増により発生することが多いので、需要の予測が難しい場合にはAWS Lambda SnapStartは有効です。

兼安 聡(執筆記事の一覧)

アプリケーションサービス部 DS3課所属
2024 Japan AWS Top Engineers (Database)
2024 Japan AWS All Certifications Engineers
Certified ScrumMaster
PMP
広島在住です。今日も明日も修行中です。