BentoML로 GPU 추론 API 배포·롤백 자동화

import bentoml
from bentoml.io import JSON

import torch
from transformers import AutoTokenizer, AutoModelForSequenceClassification

MODEL_ID = "distilbert-base-uncased-finetuned-sst-2-english"

@bentoml.service(
    name="gpu-infer-api",
    traffic={
        # 운영에서는 timeout을 충분히 주고, 프록시/Ingress 타임아웃과 맞추는 게 핵심
        "timeout": 30,
        "concurrency": 16,
    },
)
class GPUInferService:
    def __init__(self):
        self.device = "cuda" if torch.cuda.is_available() else "cpu"
        self.tokenizer = AutoTokenizer.from_pretrained(MODEL_ID)
        self.model = AutoModelForSequenceClassification.from_pretrained(MODEL_ID)
        self.model.to(self.device)
        self.model.eval()

    @bentoml.api(input=JSON(), output=JSON())
    def predict(self, payload: dict) -> dict:
        text = payload.get("text", "")
        if not text:
            return {"ok": False, "error": "text is required"}

        inputs = self.tokenizer(text, return_tensors="pt", truncation=True)
        inputs = {k: v.to(self.device) for k, v in inputs.items()}

        with torch.inference_mode():
            logits = self.model(**inputs).logits
            probs = torch.softmax(logits, dim=-1)[0].detach().cpu().tolist()

        return {"ok": True, "probs": probs, "device": self.device}

service: "service:GPUInferService"

python:
  packages:
    - bentoml==1.2.19
    - torch==2.2.2
    - transformers==4.41.2
    - accelerate==0.31.0

# 컨테이너 레벨에서 CUDA 런타임을 쓰는 경우, base image 전략을 명확히 해야 함
# (예: nvidia/cuda 기반 이미지 사용)

bentoml build
bentoml list

bentoml containerize gpu-infer-api:latest -t registry.example.com/gpu-infer-api:1.0.0

apiVersion: apps/v1
kind: Deployment
metadata:
  name: gpu-infer-api
spec:
  replicas: 2
  strategy:
    type: RollingUpdate
    rollingUpdate:
      maxSurge: 1
      maxUnavailable: 0
  selector:
    matchLabels:
      app: gpu-infer-api
  template:
    metadata:
      labels:
        app: gpu-infer-api
    spec:
      terminationGracePeriodSeconds: 60
      containers:
        - name: api
          image: registry.example.com/gpu-infer-api:1.0.0
          ports:
            - containerPort: 3000
          resources:
            limits:
              nvidia.com/gpu: 1
              cpu: "2"
              memory: "8Gi"
            requests:
              nvidia.com/gpu: 1
              cpu: "1"
              memory: "6Gi"
          readinessProbe:
            httpGet:
              path: /readyz
              port: 3000
            initialDelaySeconds: 10
            periodSeconds: 5
            timeoutSeconds: 2
            failureThreshold: 6
          livenessProbe:
            httpGet:
              path: /healthz
              port: 3000
            initialDelaySeconds: 30
            periodSeconds: 10
            timeoutSeconds: 2
            failureThreshold: 3
          startupProbe:
            httpGet:
              path: /healthz
              port: 3000
            failureThreshold: 60
            periodSeconds: 2

kubectl rollout status deployment/gpu-infer-api
kubectl rollout history deployment/gpu-infer-api
kubectl rollout undo deployment/gpu-infer-api --to-revision=3

name: deploy-gpu-infer
on:
  push:
    branches: ["main"]

jobs:
  deploy:
    runs-on: ubuntu-latest
    steps:
      - uses: actions/checkout@v4

      - name: Build and push image
        run: |
          docker build -t registry.example.com/gpu-infer-api:${GITHUB_SHA} .
          docker push registry.example.com/gpu-infer-api:${GITHUB_SHA}

      - name: Set kubeconfig
        run: |
          mkdir -p ~/.kube
          echo "${KUBECONFIG_B64}" | base64 -d > ~/.kube/config
        env:
          KUBECONFIG_B64: ${{ secrets.KUBECONFIG_B64 }}

      - name: Deploy
        run: |
          kubectl set image deployment/gpu-infer-api api=registry.example.com/gpu-infer-api:${GITHUB_SHA}

      - name: Wait rollout
        run: |
          set -e
          if ! kubectl rollout status deployment/gpu-infer-api --timeout=300s; then
            echo "rollout failed, rollback"
            kubectl rollout undo deployment/gpu-infer-api
            exit 1
          fi

apiVersion: argoproj.io/v1alpha1
kind: Rollout
metadata:
  name: gpu-infer-api
spec:
  replicas: 2
  selector:
    matchLabels:
      app: gpu-infer-api
  template:
    metadata:
      labels:
        app: gpu-infer-api
    spec:
      containers:
        - name: api
          image: registry.example.com/gpu-infer-api:1.0.0
          resources:
            limits:
              nvidia.com/gpu: 1
  strategy:
    canary:
      steps:
        - setWeight: 5
        - pause: { duration: 60 }
        - setWeight: 25
        - pause: { duration: 120 }
        - setWeight: 50
        - pause: { duration: 180 }