サーバーレス GPU の Modal でコールドスタートを計測する
はじめに
Modal というサービスをご存知だろうか。Modal はサーバーレス GPU を提供している PaaS の一つである。有名どころでは Substack が SageMaker から移行したという事例を載せており、今非常に勢いがある企業の一つだ。サーバーレス GPU というのは要は AWS Lambda の GPU 版である。
サーバーレス GPU を使いたいのはどんな時だろうか?たとえばアイテムに対する埋め込み計算やタグづけなど、アイテムが新しく追加された時だけ実行すれば良いような計算で、かつその頻度がそこまで高くないようなものは、常時 GPU サーバーを確保しておくとアイドル時間ばかりでお金が無駄になる。そういう時はバッチやイベント駆動のワークフロー中でサーバーレス GPU 上の ML モデルを呼べると実際の推論分にしか課金されず便利だろう。
サーバーレスのサービスではしばしばコールドスタートにかかる時間が問題になる。GPU を使ったモデル推論ではこの問題がさらに顕著で、GPU 用の Docker イメージは大抵サイズが非常に大きいのでインスタンスを確保してイメージを pull して立ち上げるだけで時間がかかるし、さらにその上にモデルのウェイトをロードしなければならない。汎用的なサービスでやると分単位でかかるのが普通である。
Modal などのサーバーレス GPU 業者は独自の最適化技術でこのコールドスタートを非常に短く抑えていることを売りにしている。そこで今回はいくつかのシナリオでこのコールドスタートの時間を簡単に計測してみようと思う。
Modal の使い方
Modal には training や notebook インスタンスなどいくつかのサービスがあるがここでは inference について簡単に紹介する。以下は公式ページのトップに載っているコード例である。
import modal
MODEL_NAME = "black-forest-labs/FLUX.1-schnell"
image = (
modal.Image.from_registry("nvidia/cuda:12.4.0-devel-ubuntu22.04")
.pip_install("torch", "transformers", "diffusers", ...)
)
volume = modal.Volume.from_name("flux-lora-models")
@app.cls(gpu="H100", image=image, volumes={"/loras": volume})
class FluxWithLoRA:
@modal.enter()
def setup(self):
self.pipeline = FluxPipeline.from_pretrained(MODEL_NAME).to("cuda")
self.pipeline.load_and_fuse_lora()
@modal.method()
def generate_image(self, prompt: str):
return self.pipeline(prompt).images[0]
flux = FluxWithLoRA()
flux.generate_image.remote("")
見ての通り、ベースの Docker イメージ・インストールする Python ライブラリ・GPU タイプのような設定を全て Python コードとして記述できる。@app.cls にそれらインフラ設定を書き、@modal.enter() でアプリケーションの初期化、@modal.method() で実際に呼び出す関数を実装する。デプロイもコマンドラインから modal deploy app.py を実行するだけで良い。GPU や CUDA が絡むイメージのビルドは落とし穴も多く、一回のビルドにも時間がかかるので個人的には苦痛を伴う作業だったが、Modal 側がある程度テンプレ化してくれているおかげでハマりづらい。
上記の例は SDK 経由のアクセスを想定しているが、後述のように一般的な Web エンドポイントも容易に作成可能である。
セットアップ
今回は以下の三つのセットアップで、cl-nagoya/ruri-v3-30m と cl-nagoya/ruri-v3-310m による文埋め込み計算をするエンドポイントを用意した。いずれも L4 GPU 上にデプロイし、proxy auth で認証を設定している。
Note
ビルドスクリプト・計測スクリプトは Claude Code と Codex に書いてもらった。
1. sentence-transformers + FastAPI
FastAPI 内で sentence-transformers モデルによる埋め込み計算を行なってレスポンスするシンプルな方法。一応 flash-attn も入れている。
import modal
app = modal.App("ruri-coldstart-st-cold")
model_volume = modal.Volume.from_name("ruri-weights", create_if_missing=True)
def _prefetch_30m():
from huggingface_hub import snapshot_download
snapshot_download(
repo_id="cl-nagoya/ruri-v3-30m",
local_dir="/cache/hf/models/ruri-v3-30m",
cache_dir="/cache/hf/hub",
)
model_volume.commit()
@app.cls(
image=(
modal.Image.debian_slim(python_version="3.11")
.uv_pip_install(
"torch==2.6.0",
"https://github.com/Dao-AILab/flash-attention/releases/download/v2.7.4.post1/flash_attn-2.7.4.post1+cu12torch2.6cxx11abiFALSE-cp311-cp311-linux_x86_64.whl",
"sentence-transformers>=3.3",
"transformers>=4.48",
"fastapi[standard]>=0.115",
)
.env(
{
"HF_HOME": "/cache/hf",
"HF_HUB_CACHE": "/cache/hf/hub",
"HUGGINGFACE_HUB_CACHE": "/cache/hf/hub",
"TRANSFORMERS_CACHE": "/cache/hf/transformers",
"SENTENCE_TRANSFORMERS_HOME": "/cache/hf/sentence-transformers",
"ENABLE_SNAPSHOT": "0",
}
)
.run_function(_prefetch_30m, volumes={"/cache/hf": model_volume})
),
gpu="L4",
scaledown_window=60,
max_containers=1,
volumes={"/cache/hf": model_volume},
enable_memory_snapshot=False,
)
class STRuri30m:
@modal.enter()
def load(self):
import os
os.environ["HF_HUB_OFFLINE"] = "1"
import torch
from sentence_transformers import SentenceTransformer
self.model = SentenceTransformer(
"/cache/hf/models/ruri-v3-30m",
device="cuda",
model_kwargs={
"attn_implementation": "flash_attention_2",
"torch_dtype": torch.bfloat16,
},
local_files_only=True,
trust_remote_code=True,
)
self.model.encode(["ウォームアップ"], normalize_embeddings=True)
@modal.asgi_app(requires_proxy_auth=True)
def web(self):
from fastapi import FastAPI
api = FastAPI()
@api.post("/embed")
def embed(req: dict):
inputs = req.get("inputs", [])
if not isinstance(inputs, list):
inputs = [inputs]
embs = self.model.encode(
inputs,
normalize_embeddings=True,
).tolist()
return {"embeddings": embs, "dim": len(embs[0]) if embs else 0}
return api
イメージビルド時に、_prefetch_30m でついでにモデルを Modal の Volume に保存し、実際に使う際にネットワーク越しからモデルウェイトをダウンロードするのを回避しているのがポイント。FastAPI のエンドポイントは @modal.asgi_app で作成できる。
Tip
ビルドや起動を高速化するために、モデルウェイトは Modal Volume においてコンテナにマウントするのがベストプラクティスである。
2. sentence-transformers + FastAPI + Memory Snapshots
まだアルファ版だが、Modal には Memory Snapshots という、コンテナの初期化処理をキャッシュしてくれるような機能がある。PyTorch の読み込みなどは結構時間がかかるので、これを使うと立ち上がりがより速くなるらしい。
@app.cls(
image=(
modal.Image.debian_slim(python_version="3.11")
...
.env(
{
"ENABLE_SNAPSHOT": "1",
...
}
)
.run_function(_prefetch_30m, volumes={"/cache/hf": model_volume})
),
enable_memory_snapshot=True,
experimental_options={"enable_gpu_snapshot": True},
...
)
class STRuri30m:
@modal.enter(snap=True)
def load(self):
...3. Text Embeddings Inference (TEI)
三つ目は FastAPI を立てるのではなく、Rust 製の高速な埋め込み推論サーバーである Text Embeddings Inference (TEI) を利用する。このような場合には、バックグラウンドでサーバープロセスを起動し、Modal では @modal.web_server によってリクエストをそのままプロキシするのがスタンダードなようである。1
import socket
import subprocess
import time
import modal
app = modal.App("ruri-coldstart-tei")
model_volume = modal.Volume.from_name("ruri-weights", create_if_missing=True)
PORT = 8000
def spawn_server(model_id: str) -> subprocess.Popen:
proc = subprocess.Popen(
[
"text-embeddings-router",
"--model-id",
model_id,
"--huggingface-hub-cache",
"/cache/hf/hub",
"--port",
str(PORT),
"--hostname",
"0.0.0.0",
]
)
deadline = time.monotonic() + 300
while time.monotonic() < deadline:
try:
socket.create_connection(("127.0.0.1", PORT), timeout=1).close()
return proc
except (OSError, ConnectionRefusedError):
if proc.poll() is not None:
raise RuntimeError(
f"text-embeddings-router exited with code {proc.returncode}"
)
time.sleep(0.5)
proc.terminate()
raise RuntimeError("text-embeddings-router did not become ready within 300s")
def _prefetch_30m():
spawn_server("cl-nagoya/ruri-v3-30m").terminate()
@app.cls(
image=(
modal.Image.from_registry(
"ghcr.io/huggingface/text-embeddings-inference:89-1.9.3",
add_python="3.11",
)
.dockerfile_commands("ENTRYPOINT []")
.env(
{
"HF_HOME": "/cache/hf",
"HF_HUB_CACHE": "/cache/hf/hub",
"HUGGINGFACE_HUB_CACHE": "/cache/hf/hub",
"TRANSFORMERS_CACHE": "/cache/hf/transformers",
"SENTENCE_TRANSFORMERS_HOME": "/cache/hf/sentence-transformers",
}
)
.run_function(
_prefetch_30m,
gpu="L4",
volumes={"/cache/hf": model_volume},
)
),
gpu="L4",
scaledown_window=60,
max_containers=1,
volumes={"/cache/hf": model_volume},
)
class TEIRuri30m:
@modal.enter()
def boot(self):
self.proc = spawn_server("cl-nagoya/ruri-v3-30m")
@modal.exit()
def shutdown(self):
self.proc.terminate()
@modal.web_server(port=8000, requires_proxy_auth=True, startup_timeout=300)
def serve(self):
pass計測結果
上記の三種類に大きさの異なるモデルを二種類使って計6パターンのエンドポイントを作成し、コールドスタートにかかる時間を計測してみた。なお Memory Snapshots を使う場合は事前にリクエストを送ってスナップショットが作られるようにしてから計測した。2 以下はコンテナ0の状態からリクエストを投げて返ってくるまでの時間を何回か測ったうちの一つであり、各回で数秒のぶれくらいは普通にあると思って見て頂きたい。
| パターン | モデル | コールドスタート(秒) |
|---|---|---|
| sentence-transformers | ruri-v3-30m | 16.25 |
| sentence-transformers | ruri-v3-310m | 18.29 |
| sentence-transformers (Memory Snapshotsあり) | ruri-v3-30m | 5.27 |
| sentence-transformers (Memory Snapshotsあり) | ruri-v3-310m | 8.32 |
| TEI | ruri-v3-30m | 7.38 |
| TEI | ruri-v3-310m | 7.48 |
今回の計測では、TEI と Memory Snapshots ありの sentence-transformers が同程度に速かった。TEI は Python ライブラリ読み込みのオーバーヘッドが無いため高速に起動している。Memory Snapshots が使われた場合もオーバーヘッドがなくなるので同等レベルの速さが出ていた。モデルウェイトはコンテナにマウントしたボリュームから読み込んでいるので、今回使った小型のモデルであればサイズによる差は小さかった。また GPU で推論しているので、どの場合でも埋め込み計算自体は50-100ms 程度と高速に完了していた。
さいごに
GPU コンテナが数秒で起動するというのはすごい技術だなと思う。今回の検証に使ったコードは以下。
TEI の--model-idには好きな local path を設定できるため、Volume 上のモデルを適当な場所にマウントしてそのパスを指定すればリモートからのモデルダウンロードなしに起動できる。しかし --model-idに HF 上の model_id を指定した上で、HUGGINGFACE_HUB_CACHEにマウントした Volume からモデルをロードする方が明確に起動が速かった。何かしら最適化がされているのかもしれない。
ただし、リクエストを送った先でスナップショットが利用可能になっているとは限らず(インスタンスタイプごとに必要だったりするようで、事前に作っても使われないことがある)、その場合はスナップショットの作成が走ってしまうので逆に遅くなってしまう。ここに載せた数値はスナップショットから復元が行われた場合の計測結果。