MuseTalk本番デプロイ実践 — Docker・GPUサービング・オートスケール・コスト最適化・可観測性

[API (Next.js Route Handler)]
   │  ① 入力をZod検証 → 冪等性キー → キューに投入（すぐ返す）
   ▼
[ジョブキュー (SQS / Redis Stream)]
   │
   ▼
[GPUワーカー (常駐・複数Pod)]   ← KEDAがキュー深さで台数を増減
   │  ② 起動時にモデルを1回ロードし常駐。アバターはLRUキャッシュ
   │  ③ 生成 → オブジェクトストレージへ出力 → 完了をWebhook/更新
   ▼
[オブジェクトストレージ (S3/GCS, 署名付きURL)]

# Dockerfile — 公式準拠の固定環境（CUDA 11.7 / Python 3.10 / PyTorch 2.0.1）
FROM nvidia/cuda:11.7.1-cudnn8-runtime-ubuntu22.04

ENV DEBIAN_FRONTEND=noninteractive \
    PYTHONUNBUFFERED=1 \
    PIP_NO_CACHE_DIR=1

# システム依存：ffmpeg（動画I/O）、libgl1（OpenCV）
RUN apt-get update && apt-get install -y --no-install-recommends \
      python3.10 python3-pip git ffmpeg libgl1 libglib2.0-0 \
    && rm -rf /var/lib/apt/lists/*

WORKDIR /app

# ① PyTorchはCUDA 11.7ビルドを明示（噛み合わないとCPU動作で激遅になる）
RUN pip3 install torch==2.0.1 torchvision==0.15.2 torchaudio==2.0.2 \
      --index-url https://download.pytorch.org/whl/cu117

# ② アプリ依存
COPY requirements.txt .
RUN pip3 install -r requirements.txt

# ③ mmlabは“mim”でこのバージョン固定（順序とバージョンが命）
RUN pip3 install -U openmim \
    && mim install "mmengine" \
    && mim install "mmcv==2.0.1" \
    && mim install "mmdet==3.1.0" \
    && mim install "mmpose==1.1.0"

COPY . .

# 重みはイメージに焼かない（巨大化＆Pull遅延を避ける）。起動時にキャッシュへ取得する
EXPOSE 8000
# ヘルスチェックでモデル常駐の準備完了を確認
HEALTHCHECK --interval=30s --timeout=5s --start-period=120s --retries=3 \
  CMD python3 -c "import urllib.request; urllib.request.urlopen('http://localhost:8000/healthz')" || exit 1

CMD ["python3", "-m", "uvicorn", "service.app:app", "--host", "0.0.0.0", "--port", "8000"]

# 例。実際の必要パッケージは公式リポジトリのrequirements.txtに従う
diffusers
accelerate
opencv-python
numpy
omegaconf
transformers
# fastapi / uvicorn / boto3 などサービス用も追加
fastapi
uvicorn[standard]
boto3

# service/app.py — モデル常駐＋アバターLRUキャッシュのGPU推論サービス
from contextlib import asynccontextmanager
from collections import OrderedDict
import hashlib
import torch
from fastapi import FastAPI, HTTPException
from pydantic import BaseModel, HttpUrl

# MuseTalkリポジトリの推論内部を import して薄くラップする（関数名は実装に合わせる）
from musetalk_service.engine import MuseTalkEngine  # ← 自前の薄いアダプタ層


class GenerateRequest(BaseModel):
    avatar_id: str
    source_video_url: HttpUrl
    audio_url: HttpUrl
    bbox_shift: int = 0
    use_float16: bool = True


# アバター前処理結果のLRUキャッシュ（再利用＝前処理の払い戻し）
class AvatarCache:
    def __init__(self, capacity: int = 8) -> None:
        self._store: "OrderedDict[str, object]" = OrderedDict()
        self._cap = capacity

    def get(self, key: str):
        if key in self._store:
            self._store.move_to_end(key)
            return self._store[key]
        return None

    def put(self, key: str, value: object) -> None:
        self._store[key] = value
        self._store.move_to_end(key)
        if len(self._store) > self._cap:
            self._store.popitem(last=False)  # 最も使われていないアバターを退避


@asynccontextmanager
async def lifespan(app: FastAPI):
    # 起動時に1回だけモデルをGPUへ常駐ロード（ここが重い。だから1回だけ）
    app.state.engine = MuseTalkEngine(
        device="cuda",
        dtype=torch.float16,  # fp16で省VRAM・高速
        weights_dir="/cache/models",
    )
    app.state.avatars = AvatarCache(capacity=8)
    yield
    # graceful shutdown：進行中ジョブの後始末
    app.state.engine.close()


app = FastAPI(lifespan=lifespan)


@app.get("/healthz")
def healthz():
    # モデル常駐の準備完了をK8s/Dockerのヘルスチェックへ返す
    return {"ok": app.state.engine.is_ready()}


@app.post("/generate")
def generate(req: GenerateRequest):
    key = req.avatar_id
    prepared = app.state.avatars.get(key)
    if prepared is None:
        # 新規アバターは1回だけ前処理（preparation: True 相当）
        prepared = app.state.engine.prepare(str(req.source_video_url), bbox_shift=req.bbox_shift)
        app.state.avatars.put(key, prepared)

    try:
        out_url = app.state.engine.speak(prepared, str(req.audio_url))  # 即時生成
        return {"video_url": out_url, "avatar_cached": prepared is not None}
    except torch.cuda.OutOfMemoryError:
        torch.cuda.empty_cache()
        raise HTTPException(status_code=503, detail="gpu_oom_retry")  # 上位でバッチ縮小して再試行

// app/api/lipsync/route.ts — 投入は即返し、生成はワーカーへ（Next.js Route Handler）
import { NextResponse } from "next/server";
import { createHash } from "node:crypto";
import { z } from "zod";

const Req = z.object({
  avatarId: z.string().min(1),
  sourceVideoUrl: z.string().url(),
  audioUrl: z.string().url(),
  bboxShift: z.number().int().default(0),
});

function jobKey(input: z.infer<typeof Req>): string {
  return createHash("sha256").update(JSON.stringify(input)).digest("hex");
}

export async function POST(req: Request) {
  const parsed = Req.safeParse(await req.json());
  if (!parsed.success) {
    return NextResponse.json({ error: parsed.error.flatten() }, { status: 422 });
  }
  const key = jobKey(parsed.data);

  // 冪等性：同じ入力なら作り直さない（二重生成・二重課金の防止）
  const existing = await jobStore.get(key);
  if (existing) return NextResponse.json({ jobId: existing.id, cached: true });

  const job = await queue.enqueue("lipsync", { ...parsed.data, key });
  await jobStore.put(key, { id: job.id, status: "queued" });
  return NextResponse.json({ jobId: job.id, cached: false });
}

# worker.py — 失敗を握り、回復する（少なくとも1回・冪等前提）
def handle(job: dict) -> None:
    try:
        run_generate(job, batch_size=job.get("batch_size", 8))
        mark_done(job["key"])
    except torch.cuda.OutOfMemoryError:
        torch.cuda.empty_cache()
        bs = max(1, job.get("batch_size", 8) // 2)
        requeue(job | {"batch_size": bs})  # バッチを半分にして積み直す
    except SpotInterruption:
        requeue(job)  # スポット中断は“正常”。別Podが拾う（冪等なので二重実行しても安全）

# keda-scaledobject.yaml — キュー深さでGPUワーカーをオートスケール
apiVersion: keda.sh/v1alpha1
kind: ScaledObject
metadata:
  name: musetalk-worker
spec:
  scaleTargetRef:
    name: musetalk-worker         # GPUワーカーのDeployment
  minReplicaCount: 1              # 低遅延要件：常に1台温存しコールドスタートを消す
  maxReplicaCount: 10
  cooldownPeriod: 300            # 縮小は緩やかに（生成中の打ち切りを避ける）
  triggers:
    - type: aws-sqs-queue
      metadata:
        queueURL: https://sqs.ap-northeast-1.amazonaws.com/xxx/lipsync
        queueLength: "3"         # 1台あたり3件たまったら増やす
        awsRegion: ap-northeast-1

# deployment抜粋 — GPUを要求し、graceful shutdownの猶予を取る
spec:
  template:
    spec:
      terminationGracePeriodSeconds: 120   # 生成中ジョブを取りこぼさない
      containers:
        - name: worker
          resources:
            limits:
              nvidia.com/gpu: 1

# 構造化ログ（PII除外）。相関IDで API→キュー→ワーカー→出力 を横断
log.info("lipsync_done", extra={
    "job_id": job_id,
    "avatar_id": avatar_id,
    "avatar_cache_hit": cache_hit,      # 再利用が効いたか
    "prepare_ms": prepare_ms,           # 前処理（初回のみ重い）
    "speak_ms": speak_ms,               # 生成
    "sync_score": sync_score,           # 品質ゲートのしきい値判定に使う
    "gpu_mem_peak_mb": peak_mem,
    # ❌ 入れない：音声内容・顔画像・元動画
})