AIモデルのデプロイとサービング完全ガイド：Triton、vLLM、BentoML、Kubernetes

import httpx
import asyncio

async def stream_llm_response(prompt: str):
    async with httpx.AsyncClient() as client:
        async with client.stream(
            "POST",
            "http://localhost:8000/v1/completions",
            json={
                "model": "llama-3-8b",
                "prompt": prompt,
                "max_tokens": 512,
                "stream": True
            },
            timeout=60.0
        ) as response:
            async for line in response.aiter_lines():
                if line.startswith("data: "):
                    data = line[6:]
                    if data != "[DONE]":
                        import json
                        chunk = json.loads(data)
                        token = chunk["choices"][0]["text"]
                        print(token, end="", flush=True)

# ---- ステージ1: ビルダー ----
FROM nvidia/cuda:12.1.0-cudnn8-devel-ubuntu22.04 AS builder

WORKDIR /build

RUN apt-get update && apt-get install -y \
    python3.11 \
    python3.11-dev \
    python3-pip \
    git \
    && rm -rf /var/lib/apt/lists/*

COPY requirements.txt .
RUN pip install --no-cache-dir --prefix=/install -r requirements.txt

# ---- ステージ2: ランタイム ----
FROM nvidia/cuda:12.1.0-cudnn8-runtime-ubuntu22.04

ENV PYTHONUNBUFFERED=1
ENV NVIDIA_VISIBLE_DEVICES=all
ENV NVIDIA_DRIVER_CAPABILITIES=compute,utility

WORKDIR /app

RUN apt-get update && apt-get install -y \
    python3.11 \
    python3-pip \
    libgomp1 \
    && rm -rf /var/lib/apt/lists/*

# ビルダーステージからインストール済みパッケージのみコピー
COPY --from=builder /install /usr/local
COPY . .

# セキュリティのため非rootユーザーで実行
RUN useradd -m -u 1000 mluser
USER mluser

EXPOSE 8080
CMD ["python3.11", "-m", "uvicorn", "app.main:app", "--host", "0.0.0.0", "--port", "8080"]

# UbuntuへのNVIDIA Container Toolkitのインストール
distribution=$(. /etc/os-release;echo $ID$VERSION_ID)
curl -fsSL https://nvidia.github.io/libnvidia-container/gpgkey | sudo gpg --dearmor \
  -o /usr/share/keyrings/nvidia-container-toolkit-keyring.gpg

curl -s -L https://nvidia.github.io/libnvidia-container/$distribution/libnvidia-container.list | \
  sed 's#deb https://#deb [signed-by=/usr/share/keyrings/nvidia-container-toolkit-keyring.gpg] https://#g' | \
  sudo tee /etc/apt/sources.list.d/nvidia-container-toolkit.list

sudo apt-get update
sudo apt-get install -y nvidia-container-toolkit
sudo nvidia-ctk runtime configure --runtime=docker
sudo systemctl restart docker

# GPUコンテナアクセスの確認
docker run --rm --gpus all nvidia/cuda:12.1.0-base-ubuntu22.04 nvidia-smi

apiVersion: apps/v1
kind: Deployment
metadata:
  name: model-serving
  namespace: ml-serving
spec:
  replicas: 2
  selector:
    matchLabels:
      app: model-serving
  template:
    metadata:
      labels:
        app: model-serving
      annotations:
        prometheus.io/scrape: 'true'
        prometheus.io/port: '8080'
        prometheus.io/path: '/metrics'
    spec:
      nodeSelector:
        accelerator: nvidia-tesla-a10g
      tolerations:
        - key: nvidia.com/gpu
          operator: Exists
          effect: NoSchedule
      containers:
        - name: model-server
          image: myregistry/model-server:v1.2.0
          ports:
            - containerPort: 8080
          resources:
            requests:
              cpu: '2'
              memory: '8Gi'
              nvidia.com/gpu: '1'
            limits:
              cpu: '4'
              memory: '16Gi'
              nvidia.com/gpu: '1'
          env:
            - name: MODEL_PATH
              value: '/models/llama-3-8b'
            - name: MAX_BATCH_SIZE
              value: '32'
          livenessProbe:
            httpGet:
              path: /health
              port: 8080
            initialDelaySeconds: 60
            periodSeconds: 15
          readinessProbe:
            httpGet:
              path: /ready
              port: 8080
            initialDelaySeconds: 30
            periodSeconds: 10
---
apiVersion: autoscaling/v2
kind: HorizontalPodAutoscaler
metadata:
  name: model-serving-hpa
  namespace: ml-serving
spec:
  scaleTargetRef:
    apiVersion: apps/v1
    kind: Deployment
    name: model-serving
  minReplicas: 2
  maxReplicas: 10
  metrics:
    - type: Resource
      resource:
        name: cpu
        target:
          type: Utilization
          averageUtilization: 60
    - type: Pods
      pods:
        metric:
          name: model_requests_per_second
        target:
          type: AverageValue
          averageValue: '50'

apiVersion: karpenter.sh/v1beta1
kind: NodePool
metadata:
  name: gpu-nodepool
spec:
  template:
    metadata:
      labels:
        accelerator: nvidia-tesla-a10g
    spec:
      nodeClassRef:
        name: gpu-nodeclass
      requirements:
        - key: karpenter.sh/capacity-type
          operator: In
          values: ['on-demand', 'spot']
        - key: node.kubernetes.io/instance-type
          operator: In
          values: ['g5.xlarge', 'g5.2xlarge', 'g5.4xlarge']
        - key: kubernetes.io/arch
          operator: In
          values: ['amd64']
      taints:
        - key: nvidia.com/gpu
          effect: NoSchedule
  limits:
    nvidia.com/gpu: '20'
  disruption:
    consolidationPolicy: WhenUnderutilized
    consolidateAfter: 30s

# config.pbtxt
name: "text_classifier"
platform: "onnxruntime_onnx"
max_batch_size: 64

input [
  {
    name: "input_ids"
    data_type: TYPE_INT64
    dims: [ -1 ]
  },
  {
    name: "attention_mask"
    data_type: TYPE_INT64
    dims: [ -1 ]
  }
]

output [
  {
    name: "logits"
    data_type: TYPE_FP32
    dims: [ -1, 2 ]
  }
]

dynamic_batching {
  preferred_batch_size: [ 8, 16, 32 ]
  max_queue_delay_microseconds: 5000
}

instance_group [
  {
    count: 2
    kind: KIND_GPU
    gpus: [ 0 ]
  }
]

import bentoml
from bentoml.io import JSON
from pydantic import BaseModel
import numpy as np
from typing import List

class InferenceRequest(BaseModel):
    texts: List[str]
    top_k: int = 5

class InferenceResponse(BaseModel):
    labels: List[str]
    scores: List[float]

# モデルランナーの作成
classifier_runner = bentoml.pytorch.get("text_classifier:latest").to_runner()

svc = bentoml.Service("text_classification_svc", runners=[classifier_runner])

@svc.api(input=JSON(pydantic_model=InferenceRequest),
         output=JSON(pydantic_model=InferenceResponse))
async def classify(request: InferenceRequest) -> InferenceResponse:
    batch_results = await classifier_runner.async_run(request.texts)

    labels = []
    scores = []
    for result in batch_results:
        top_idx = np.argsort(result)[-request.top_k:][::-1]
        labels.extend([f"label_{i}" for i in top_idx])
        scores.extend(result[top_idx].tolist())

    return InferenceResponse(labels=labels, scores=scores)

from vllm import LLM, SamplingParams
from vllm.engine.arg_utils import AsyncEngineArgs
from vllm.engine.async_llm_engine import AsyncLLMEngine
from fastapi import FastAPI
from fastapi.responses import StreamingResponse
import uvicorn
import json
import uuid

app = FastAPI(title="vLLM OpenAI互換API")

engine_args = AsyncEngineArgs(
    model="meta-llama/Llama-3-8B-Instruct",
    tensor_parallel_size=2,       # 2-GPUテンソル並列処理
    gpu_memory_utilization=0.90,
    max_model_len=8192,
    enable_chunked_prefill=True,
)
engine = AsyncLLMEngine.from_engine_args(engine_args)

@app.post("/v1/chat/completions")
async def chat_completions(request: dict):
    messages = request.get("messages", [])
    stream = request.get("stream", False)

    prompt = format_messages(messages)
    sampling_params = SamplingParams(
        temperature=request.get("temperature", 0.7),
        max_tokens=request.get("max_tokens", 512),
        top_p=request.get("top_p", 0.95),
    )

    request_id = str(uuid.uuid4())

    if stream:
        async def generate_stream():
            async for output in engine.generate(prompt, sampling_params, request_id):
                if output.outputs:
                    token = output.outputs[0].text
                    chunk = {
                        "id": request_id,
                        "object": "chat.completion.chunk",
                        "choices": [{"delta": {"content": token}, "index": 0}]
                    }
                    yield f"data: {json.dumps(chunk)}\n\n"
            yield "data: [DONE]\n\n"

        return StreamingResponse(generate_stream(), media_type="text/event-stream")

    final_output = None
    async for output in engine.generate(prompt, sampling_params, request_id):
        final_output = output

    return {
        "choices": [{"message": {"content": final_output.outputs[0].text}}]
    }

def format_messages(messages):
    result = ""
    for msg in messages:
        role = msg.get("role", "user")
        content = msg.get("content", "")
        result += f"<|{role}|>\n{content}\n"
    return result + "<|assistant|>\n"

# TGIでLLaMA-3サービングを開始
docker run --gpus all \
  -p 8080:80 \
  -v /data/models:/data \
  -e HUGGING_FACE_HUB_TOKEN=$HF_TOKEN \
  ghcr.io/huggingface/text-generation-inference:2.0.4 \
  --model-id meta-llama/Llama-3-8B-Instruct \
  --num-shard 2 \
  --max-input-length 4096 \
  --max-total-tokens 8192 \
  --max-batch-prefill-tokens 16384 \
  --dtype bfloat16

# ストリーミングレスポンスのテスト
curl http://localhost:8080/generate_stream \
  -H 'Content-Type: application/json' \
  -d '{"inputs": "What is the capital of France?", "parameters": {"max_new_tokens": 100, "stream": true}}'

import asyncio
import httpx
import logging

logger = logging.getLogger(__name__)

async def warmup_model(base_url: str, num_warmup_requests: int = 5):
    """サーバー起動時にモデルウォームアップを実行します。"""
    dummy_request = {
        "inputs": "warmup",
        "parameters": {"max_new_tokens": 10}
    }

    async with httpx.AsyncClient(timeout=120.0) as client:
        logger.info("モデルウォームアップを開始します...")

        # ヘルスチェックの待機
        for _ in range(30):
            try:
                resp = await client.get(f"{base_url}/health")
                if resp.status_code == 200:
                    break
            except Exception:
                await asyncio.sleep(2)

        # ウォームアップリクエストの送信
        tasks = [
            client.post(f"{base_url}/generate", json=dummy_request)
            for _ in range(num_warmup_requests)
        ]
        await asyncio.gather(*tasks, return_exceptions=True)
        logger.info(f"ウォームアップ完了（{num_warmup_requests}件のリクエストを送信）")

import asyncio
from collections import deque
from dataclasses import dataclass, field
from typing import Any

@dataclass
class BatchRequest:
    request_id: str
    payload: dict
    future: asyncio.Future = field(default_factory=asyncio.Future)

class DynamicBatcher:
    def __init__(self, max_batch_size: int = 32, max_wait_ms: float = 10.0):
        self.max_batch_size = max_batch_size
        self.max_wait_ms = max_wait_ms
        self.queue: deque = deque()
        self._lock = asyncio.Lock()

    async def add_request(self, request_id: str, payload: dict) -> Any:
        req = BatchRequest(request_id=request_id, payload=payload)
        async with self._lock:
            self.queue.append(req)
        return await req.future

    async def process_batches(self, model_fn):
        while True:
            await asyncio.sleep(self.max_wait_ms / 1000)
            async with self._lock:
                if not self.queue:
                    continue
                batch = []
                while self.queue and len(batch) < self.max_batch_size:
                    batch.append(self.queue.popleft())

            if batch:
                try:
                    inputs = [r.payload for r in batch]
                    results = await model_fn(inputs)
                    for req, result in zip(batch, results):
                        req.future.set_result(result)
                except Exception as e:
                    for req in batch:
                        req.future.set_exception(e)

from prometheus_client import Histogram, Counter, Gauge, generate_latest
from fastapi import FastAPI, Request, Response
import time

REQUEST_LATENCY = Histogram(
    "model_request_latency_seconds",
    "モデル推論のレイテンシ（秒）",
    ["model_name", "endpoint"],
    buckets=[0.01, 0.05, 0.1, 0.25, 0.5, 1.0, 2.5, 5.0, 10.0]
)

REQUEST_COUNT = Counter(
    "model_requests_total",
    "推論リクエストの総数",
    ["model_name", "endpoint", "status"]
)

ACTIVE_REQUESTS = Gauge(
    "model_active_requests",
    "現在処理中のリクエスト数",
    ["model_name"]
)

TOKEN_THROUGHPUT = Counter(
    "model_tokens_generated_total",
    "生成されたトークンの総数",
    ["model_name"]
)

GPU_MEMORY_USED = Gauge(
    "gpu_memory_used_bytes",
    "GPU使用メモリ（バイト）",
    ["gpu_index"]
)

app = FastAPI()

@app.middleware("http")
async def metrics_middleware(request: Request, call_next):
    model_name = "llama-3-8b"
    endpoint = request.url.path

    ACTIVE_REQUESTS.labels(model_name=model_name).inc()
    start = time.perf_counter()

    try:
        response = await call_next(request)
        status = str(response.status_code)
        REQUEST_COUNT.labels(
            model_name=model_name, endpoint=endpoint, status=status
        ).inc()
        return response
    except Exception as e:
        REQUEST_COUNT.labels(
            model_name=model_name, endpoint=endpoint, status="500"
        ).inc()
        raise
    finally:
        latency = time.perf_counter() - start
        REQUEST_LATENCY.labels(
            model_name=model_name, endpoint=endpoint
        ).observe(latency)
        ACTIVE_REQUESTS.labels(model_name=model_name).dec()

@app.get("/metrics")
async def metrics():
    return Response(generate_latest(), media_type="text/plain")