FastAPI로 AI 추론 서버 구축하기 — LLM API 서빙 실전 가이드

왜 FastAPI인가

AI 추론 서버에 FastAPI가 선택받는 이유:

• 비동기 I/O: 추론 대기 중 다른 요청 처리 가능
• 타입 힌트 + Pydantic: 요청/응답 스키마 자동 검증
• 자동 API 문서: /docs Swagger UI 기본 제공
• Python 생태계: PyTorch, HuggingFace와 자연스럽게 통합

기본 추론 서버 구조

# main.py
from fastapi import FastAPI, HTTPException
from pydantic import BaseModel
from contextlib import asynccontextmanager
import torch
from transformers import AutoTokenizer, AutoModelForCausalLM

model = None
tokenizer = None

@asynccontextmanager
async def lifespan(app: FastAPI):
    # 서버 시작 시 모델 로드 (1회)
    global model, tokenizer
    tokenizer = AutoTokenizer.from_pretrained("your-model")
    model = AutoModelForCausalLM.from_pretrained(
        "your-model",
        torch_dtype=torch.float16,
        device_map="auto",
    )
    model.eval()
    yield
    # 서버 종료 시 정리
    del model, tokenizer

app = FastAPI(lifespan=lifespan)

class InferRequest(BaseModel):
    prompt: str
    max_tokens: int = 256
    temperature: float = 0.7

class InferResponse(BaseModel):
    text: str
    tokens_used: int

@app.post("/infer", response_model=InferResponse)
async def infer(req: InferRequest):
    inputs = tokenizer(req.prompt, return_tensors="pt").to(model.device)
    with torch.no_grad():
        output = model.generate(
            **inputs,
            max_new_tokens=req.max_tokens,
            temperature=req.temperature,
            do_sample=True,
        )
    text = tokenizer.decode(output[0], skip_special_tokens=True)
    return InferResponse(text=text, tokens_used=len(output[0]))

배치 추론으로 처리량 향상

단일 요청씩 처리하면 GPU 활용률이 낮습니다. 배치로 묶어 처리하세요.

import asyncio
from collections import deque

batch_queue = deque()
BATCH_SIZE = 8
BATCH_TIMEOUT = 0.05  # 50ms

async def batch_processor():
    while True:
        await asyncio.sleep(BATCH_TIMEOUT)
        if not batch_queue:
            continue
        batch = []
        while batch_queue and len(batch) < BATCH_SIZE:
            batch.append(batch_queue.popleft())
        
        # 배치 추론
        prompts = [item["prompt"] for item in batch]
        inputs = tokenizer(prompts, return_tensors="pt", padding=True).to(model.device)
        with torch.no_grad():
            outputs = model.generate(**inputs, max_new_tokens=256)
        
        for i, item in enumerate(batch):
            text = tokenizer.decode(outputs[i], skip_special_tokens=True)
            item["future"].set_result(text)

@app.on_event("startup")
async def start_batch_processor():
    asyncio.create_task(batch_processor())

타임아웃과 에러 처리

import asyncio
from fastapi import HTTPException

@app.post("/infer")
async def infer(req: InferRequest):
    try:
        result = await asyncio.wait_for(
            run_inference(req),
            timeout=30.0  # 30초 타임아웃
        )
        return result
    except asyncio.TimeoutError:
        raise HTTPException(status_code=504, detail="Inference timeout")
    except torch.cuda.OutOfMemoryError:
        torch.cuda.empty_cache()
        raise HTTPException(status_code=503, detail="GPU OOM, retry later")

헬스체크 엔드포인트

@app.get("/health")
async def health():
    gpu_available = torch.cuda.is_available()
    gpu_memory = {}
    if gpu_available:
        gpu_memory = {
            "allocated": f"{torch.cuda.memory_allocated() / 1e9:.2f}GB",
            "reserved": f"{torch.cuda.memory_reserved() / 1e9:.2f}GB",
        }
    return {
        "status": "ok",
        "model_loaded": model is not None,
        "gpu": gpu_available,
        "gpu_memory": gpu_memory,
    }

Kubernetes 배포 설정

apiVersion: apps/v1
kind: Deployment
metadata:
  name: ai-server
spec:
  replicas: 2
  template:
    spec:
      containers:
      - name: ai-server
        image: your-registry/ai-server:latest
        resources:
          limits:
            nvidia.com/gpu: 1
            memory: "16Gi"
          requests:
            nvidia.com/gpu: 1
            memory: "12Gi"
        livenessProbe:
          httpGet:
            path: /health
            port: 8000
          initialDelaySeconds: 120  # 모델 로딩 시간 확보
          periodSeconds: 30

실행

# 개발
uvicorn main:app --reload --host 0.0.0.0 --port 8000

# 프로덕션 (workers=1 권장 - GPU 메모리 공유 문제)
uvicorn main:app --host 0.0.0.0 --port 8000 --workers 1

# 또는 gunicorn + uvicorn worker
gunicorn main:app -w 1 -k uvicorn.workers.UvicornWorker --bind 0.0.0.0:8000

프로덕션에서 workers > 1은 주의: 각 worker가 GPU 메모리를 별도로 점유합니다.
수평 확장은 Pod 복제로 하고, 각 Pod는 workers=1로 유지하세요.