AWS Lambdaを高速に立ち上げるSnapStartの機能がPythonと .NETでも利用可能になりました

Lambda SnapStartの仕組み

AWS LambdaはFirecrakerというAWSが開発したサーバーレスとコンテナ専用の仮想環境上で動いています。Lambda SnapStartが有効になっていると最初に実行される初期化処理のメモリやディスクの状態をスナップショットとして仮想環境上に暗号化して保存します。そして、それ以降実行されたLambdaファンクションはゼロから初期化処理を始めるのではなく、スナップショットから新たに実行環境を生成することで高速なファンクションの立ち上げを実現します。また、スナップショットはAWS Lambdaのバージョンとエイリアスに紐づいているため、関数を更新するたびに新たなスナップショットが生成されます。（Lambdaのウォームアップとはまた別の話なので混乱することなく切り分けして捉えましましょう）

特にPythonのケースではMLの分野でSnapStartの機能が有用になるのではないでしょうか。Numpy, Pandasといった機械学習系のライブラリの読み込み、Djangoのようなフレームワークの読み込みなど、どうしてもPythonを使って機械学習系の処理を行なうと依存関係のサイズは大きくなってしまいます。また、サイズの大きなMLモデルをLmabdaで呼び出して推論を行うといったケースもあるでしょう。このようなケースでは依存関係の読み込みだけでも数十秒かかってしまうケースもあります。Lambda SnapStartを使用することでこれらのパフォーマンスを大きく向上させることが出来そうです。

.NETの場合では多くのケースで恩恵を受けることが出来るでしょう。.NETでは、JIT（Just-In-Time: ジャストインタイム）コンパイルという、プログラムの実行時に必要なコードをその場でコンパイルする方式が提供されています。しかし、このコンパイルが数十秒かかるケースがあり、AWS Lambdaに.NETを採用する一つの障壁となっていました。しかし、Lambda SnapStartを使用することでこの時間を大きく短縮することが出来るようになりました。

C#.NETでLambda SnapStartを試す

本記事ではC# .NETでLambda SnapStartを試して、どのくらいパフォーマンスが改善されるかを見てみてたいと思います。Amazon API Gatewy + AWS Lambdaの構成を作り、SnapStartを有効化したSAMテンプレートを使ってデプロイします。

Resources:
  SnapStartFunction:
    Type: AWS::Serverless::Function
    Properties:
      Handler: SnapStart::SnapStartFunctionHandler::FunctionHandler
      Runtime: dotnet8
      CodeUri: ./src/SnapStart
      MemorySize: 512
      Timeout: 15
      Policies:
        - AWSLambdaBasicExecutionRole
      Events:
        ApiGateway:
          Type: Api
          Properties:
            Path: /hello
            Method: GET
      SnapStart:
        ApplyOn: PublishedVersions
      AutoPublishAlias: SnapStart

デプロイをすると以下のようにLambda Snap Startが有効になっていることが分かります

C#のコードは以下のように記述しました。初期化処理に時間がかかるようにLambdaハンドラーのコンストラクタで大量の計算をさせてメモリに保存。関数のInvoke時にその値をAPIレスポンスに返すようにしています。通常のコールドスタートの場合はこの計算処理に数秒の時間がかかります。しかし、LambdaのSnapStartが動いていれば、計算結果はスナップショットの中に保存されているため、すぐにレスポンスが返ってくるはずです。

また、ここで注目しておきたいのが、ランタイムフックの機能です。Amazon.Lambda.Core package (version 2.5 or later)から提供されているRegisterBeforeSnapshot() メソッドはLambdaがスナップショットを作成する前に実行するコードを登録します。さらにRegisterAfterRestore() メソッドはLambdaがスナップショットから関数を再開したときに実行するコードを登録します。

これらのフックを使うことで、スナップショット登録時に動的に設定値を変更したり、外部サービスやシステムと統合して、通知の送信や外部サービスのステータスの更新などをキャッシュと両立させながら行なうことが可能になります。

using System;
using System.Collections.Generic;
using System.Threading;
using System.Threading.Tasks;
using Amazon.Lambda.APIGatewayEvents;
using Amazon.Lambda.Core;

[assembly: LambdaSerializer(typeof(Amazon.Lambda.Serialization.SystemTextJson.DefaultLambdaJsonSerializer))]

public class SnapStartFunctionHandler
{
    private static double _computationResult;
    static SnapStartFunctionHandler()
    {
        Console.WriteLine("Performing heavy computation before snapshot...");
        _computationResult = 0; // 初期化
        for (int i = 0; i < 1_000_000_000; i++) // スナップショット前に計算処理を追加
        {
            _computationResult += Math.Sqrt(i);
        }
        Console.WriteLine($"Computation result before snapshot: {_computationResult}");
    }
    public SnapStartFunctionHandler()
    {
        Amazon.Lambda.Core.SnapshotRestore.RegisterBeforeSnapshot(BeforeCheckpoint);
        Amazon.Lambda.Core.SnapshotRestore.RegisterAfterRestore(AfterCheckpoint);
    }

    private ValueTask BeforeCheckpoint()
    {
        // Add logic to be executed before taking the snapshot
        return ValueTask.CompletedTask;
    }

    private ValueTask AfterCheckpoint()
    {
        // Add logic to be executed after restoring the snapshot
        Console.WriteLine("Snap Start overloaded");
        return ValueTask.CompletedTask;
    }

    public APIGatewayProxyResponse FunctionHandler(APIGatewayProxyRequest request, ILambdaContext context)
    {
        // Add business logic

        return new APIGatewayProxyResponse
        {
            StatusCode = 200,
            Body = _computationResult.ToString()
        };
    }
}

Lambda SnapStartを無効化している場合

API Gatewayにリクエストを送ります。計算結果がレスポンスとして返ってきました。

 % curl https://6qsek89582.execute-api.us-east-1.amazonaws.com/Prod/hello
21081851051977.78                        

その結果をX-Rayで見てみると以下のようになっています。AWS Lambdaのinitフェーズでは計算処理のために4.82秒もかかっています

Lambda SnapStartを有効化している場合

API Gatewayにリクエストを送ります。当然同じ計算結果が返ってきます。

 % curl https://lwidkzdbcd.execute-api.us-east-1.amazonaws.com/Prod/hello
21081851051977.78  

X-Rayでの結果は以下のとおりです。restoreとなっている箇所がスナップショットからの復元処理ですが、532ミリ秒と大きくパフォーマンス改善されていることが分かります。

Lambda SnapStartのコスト

JavaのSnapStartは無料で使えますが、Pythonと.NETは課金がされるので注意が必要です。現状、SnapStartが紐づいているLambdaファンクションのバージョンが存在していることによる課金と、スナップショットから復元が実行されるたびに課金がされます。詳しくはAWS Lambdaの課金ページを参照してください。

まとめ

やはり、Pythonと来たら筆者としてはNode.jsも早く対応してほしいなと思いますね。特にフロントエンドなどユーザが扱う部分と同期的に処理を行う箇所でコールドスタートを気にしなくて良くなるというのは嬉しいです。以前、AWS Lambdaをモノリシックに使うLambdalithアプローチの紹介とそのメリットという記事を書きましたが、Lambdalith構成の場合はSnapStartとの相性もよく効果的に使えるのではないかなと思います。