Lambda MicroVMsのライフサイクルフックで、MLモデル推論の初回待ち時間を短縮できるか試してみた

import os
import time
import psutil
from flask import Flask, request, jsonify
from sentence_transformers import SentenceTransformer

app = Flask(__name__)

MODEL_NAME = os.environ.get("MODEL_NAME", "intfloat/multilingual-e5-small")
MODEL_DIR = os.environ.get("MODEL_DIR", "/opt/model")

# グローバルスコープでモデルロード
_start = time.time()
model = SentenceTransformer(MODEL_DIR)
MODEL_LOAD_TIME = time.time() - _start

@app.route("/")
def health():
    mem = psutil.Process().memory_info().rss / 1024 / 1024
    return jsonify(
        model_name=MODEL_NAME,
        model_load_time_sec=round(MODEL_LOAD_TIME, 3),
        memory_usage_mb=round(mem, 1),
    )

@app.route("/embed", methods=["POST"])
def embed():
    data = request.get_json(force=True)
    text = data.get("text", "hello world")
    start = time.time()
    vec = model.encode([text])
    elapsed = time.time() - start
    return jsonify(
        elapsed_sec=round(elapsed, 4),
        shape=list(vec.shape),
    )

@app.route("/aws/lambda-microvms/runtime/v1/ready", methods=["GET", "POST"])
def ready():
    """ビルド時フック: モデルロード完了+ウォームアップ後に200"""
    model.encode(["warmup"])
    return jsonify(status="ready"), 200

@app.route("/aws/lambda-microvms/runtime/v1/validate", methods=["GET", "POST"])
def validate():
    """ビルド時フック: スナップショット復元検証+模擬推論でprefetchサンプリング"""
    model.encode(["validate prefetch sampling"])
    return jsonify(status="validated"), 200

FROM public.ecr.aws/lambda/microvms:al2023-minimal

ARG MODEL_NAME=intfloat/multilingual-e5-small

RUN dnf install -y python3.11 python3.11-pip && dnf clean all && \
    ln -sf /usr/bin/python3.11 /usr/bin/python3 && \
    ln -sf /usr/bin/pip3.11 /usr/bin/pip3

COPY requirements.txt /opt/app/requirements.txt
RUN pip3 install --no-cache-dir torch --index-url https://download.pytorch.org/whl/cpu && \
    pip3 install --no-cache-dir -r /opt/app/requirements.txt psutil

# モデルを事前ダウンロード
RUN python3 -c "from sentence_transformers import SentenceTransformer; SentenceTransformer('${MODEL_NAME}').save('/opt/model')"

ENV MODEL_NAME=${MODEL_NAME}
ENV MODEL_DIR=/opt/model

COPY app.py /opt/app/app.py

WORKDIR /opt/app
CMD ["gunicorn", "app:app", "--bind", "0.0.0.0:8080", "--bind", "0.0.0.0:9000", "--timeout", "120", "--workers", "1", "--preload"]

flask==3.1.1
gunicorn==23.0.0
sentence-transformers>=3.0.0

aws lambda create-microvm-image \
  --image-name e5-small-no-hooks \
  --base-image-arn arn:aws:lambda:ap-northeast-1:aws:microvm-image:al2023-1 \
  --source-uri s3://microvm-webshell-123456789012/lifecycle-hooks/e5-small-no-hooks.zip \
  --role-arn arn:aws:iam::123456789012:role/MicroVMBuildRole

# ヘルスチェック（初回リクエスト = モデルロード発生）
$ curl -w "\nTIME: %{time_total}s\n" https://<endpoint>.lambda-microvm.ap-northeast-1.on.aws/
{"memory_usage_mb":726.0,"model_load_time_sec":18.438,"model_name":"intfloat/multilingual-e5-small"}
TIME: 18.171s

# 初回推論
$ curl -X POST -d '{"text":"Lambda MicroVMs lifecycle hooks test"}' \
  https://<endpoint>.lambda-microvm.ap-northeast-1.on.aws/embed
{"elapsed_sec":12.9989,"shape":[1,384]}

# 2回目
$ curl -X POST -d '{"text":"Lambda MicroVMs lifecycle hooks test"}' \
  https://<endpoint>.lambda-microvm.ap-northeast-1.on.aws/embed
{"elapsed_sec":0.0217,"shape":[1,384]}

aws lambda create-microvm-image \
  --image-name e5-small-ready \
  --base-image-arn arn:aws:lambda:ap-northeast-1:aws:microvm-image:al2023-1 \
  --source-uri s3://microvm-webshell-123456789012/lifecycle-hooks/e5-small-ready.zip \
  --role-arn arn:aws:iam::123456789012:role/MicroVMBuildRole \
  --hooks '{"port":9000,"microvmImageHooks":{"ready":"ENABLED","readyTimeoutInSeconds":300}}'

# ヘルスチェック
$ curl -w "\nTIME: %{time_total}s\n" https://<endpoint>.lambda-microvm.ap-northeast-1.on.aws/
{"memory_usage_mb":762.2,"model_load_time_sec":4.674,"model_name":"intfloat/multilingual-e5-small"}
TIME: 2.087s

# 初回推論
$ curl -X POST -d '{"text":"Lambda MicroVMs lifecycle hooks test"}' \
  https://<endpoint>.lambda-microvm.ap-northeast-1.on.aws/embed
{"elapsed_sec":3.4128,"shape":[1,384]}

aws lambda create-microvm-image \
  --image-name e5-small-ready-validate \
  --base-image-arn arn:aws:lambda:ap-northeast-1:aws:microvm-image:al2023-1 \
  --source-uri s3://microvm-webshell-123456789012/lifecycle-hooks/e5-small-ready-validate.zip \
  --role-arn arn:aws:iam::123456789012:role/MicroVMBuildRole \
  --hooks '{"port":9000,"microvmImageHooks":{"ready":"ENABLED","readyTimeoutInSeconds":300,"validate":"ENABLED","validateTimeoutInSeconds":300}}'

# ヘルスチェック
$ curl -w "\nTIME: %{time_total}s\n" https://<endpoint>.lambda-microvm.ap-northeast-1.on.aws/
{"memory_usage_mb":761.5,"model_load_time_sec":1.929,"model_name":"intfloat/multilingual-e5-small"}
TIME: 2.082s

# 初回推論
$ curl -X POST -d '{"text":"Lambda MicroVMs lifecycle hooks test"}' \
  https://<endpoint>.lambda-microvm.ap-northeast-1.on.aws/embed
{"elapsed_sec":1.2885,"shape":[1,384]}