MLX를 FastAPI 서비스로 감싸기

내장 서버를 넘어설 때

lm.lesson6에서 다룬 mlx-lm의 내장 mlx_lm.server는 개발과 시연, 사용자 한 명, 동시 요청이 적은 환경에 좋아. 하지만 실제 동시 요청 처리와 요청 대기열, 인증, 맞춤 기록처럼 ‘OpenAI 형식 요청에 모델이 답한다’는 수준을 넘어선 운영 요구가 생기면 직접 제어하는 FastAPI 서비스로 mlx-lm을 감싸는 편이 나아.

거창한 래퍼가 필요한 건 아니야. 시작할 때 모델을 한 번 불러오고 /generate와 /health를 공개하는 FastAPI 코드 약 100줄이면 돼. 내장 서버의 기본값과 씨름하지 않고 운영 결정을 직접 내리며 각 부분을 바꾸거나 갈라 확장할 수 있다는 게 장점이야.

기본 구조

시작할 때 모델을 한 번만 불러와 — FastAPI의 lifespan 문맥 관리자가 맡아. (model, tokenizer) 쌍을 앱 상태에 두고 모든 요청에서 다시 써.
미리 예열해 — 시작 과정에서 버릴 순전파를 한 번 실행해. 새로 불러온 뒤 첫 추론은 MLX의 JIT 컴파일 비용을 내므로, 예열하면 첫 실제 요청이 그 비용을 떠안지 않아.
/generate를 공개해 — 프롬프트와 샘플링 설정을 담은 JSON 본문을 받아 생성된 텍스트를 돌려줘. 토큰마다 보내려면 Server-Sent Events 스트리밍을 더해.
/health를 공개해 — 부하 분산기나 프로세스 관리자가 서비스 생존 여부를 싸게 확인할 수 있는 경로야.
uvicorn으로 실행해 — MLX에는 프로세스 하나, 작업자 하나를 써. 모델 생성 호출은 명시적인 잠금이나 대기열로 한 번에 하나씩 보내고, 동기 generate는 동기 엔드포인트나 스레드 오프로딩으로 이벤트 루프 밖에서 실행해.

피해야 할 것

같은 모델에 여러 작업자 프로세스를 띄우지 마. 각 프로세스가 모델 사본을 따로 불러와 GPU 메모리를 놓고 다투게 돼. 하지만 FastAPI의 async만 붙인다고 안전한 생성 동시성이 생기지도 않아. 위 최소 예제처럼 동기 generate를 비동기 경로에서 바로 호출하면 이벤트 루프가 막혀. 생성 호출을 명시적인 잠금이나 대기열로 직렬화하고 이벤트 루프 밖으로 넘겨. 운영 트래픽에서는 대기열 길이 제한과 역압, 과부하 거절 정책도 따로 정해야 해.

최소 FastAPI 서비스

아래 코드 블록은 요청 하나씩 처리하는 파일 하나짜리 최소 서비스야. app.py로 저장하고 uvicorn app:app --host 0.0.0.0 --port 8000으로 실행한 뒤, JSON 본문과 함께 POST /generate를 호출해. 여러 요청을 받으려면 앞에서 설명한 직렬화와 오프로딩을 먼저 더하고, 속도 제한과 인증, 구조화된 기록 같은 운영 기능은 그다음 붙여.

Code

MLX용 최소 FastAPI 서비스(app.py)·python

# Save as app.py; run with: uvicorn app:app --host 0.0.0.0 --port 8000

from contextlib import asynccontextmanager
from fastapi import FastAPI
from pydantic import BaseModel
from mlx_lm import load, generate

# Module-level cache for the loaded (model, tokenizer) pair
state: dict = {}


@asynccontextmanager
async def lifespan(app: FastAPI):
    # Startup: load the model once
    print("Loading model...")
    state["model"], state["tokenizer"] = load("mlx-community/Llama-3.2-1B-Instruct-4bit")

    # Warm up — pay the JIT cost during startup, not during the first real request
    _ = generate(state["model"], state["tokenizer"], prompt="warmup", max_tokens=1)
    print("Model loaded and warm.")

    yield   # App is alive here

    # Shutdown: nothing special needed; MLX cleans up at process exit
    state.clear()


app = FastAPI(lifespan=lifespan)


class GenerateRequest(BaseModel):
    prompt: str
    max_tokens: int = 100


@app.get("/health")
async def health():
    return {"status": "ok", "model_loaded": "model" in state}


@app.post("/generate")
async def gen(req: GenerateRequest):
    text = generate(
        state["model"], state["tokenizer"],
        prompt=req.prompt, max_tokens=req.max_tokens, verbose=False,
    )
    return {"text": text}

curl로 호출하기·bash

# Health check
curl http://localhost:8000/health
# {"status":"ok","model_loaded":true}

# Generate
curl -X POST http://localhost:8000/generate \
  -H "Content-Type: application/json" \
  -d '{"prompt": "Capital of France?", "max_tokens": 20}'
# {"text":"Paris is the capital of France..."}

Exercise

코드 블록의 FastAPI 서비스를 app.py로 저장해. 한 터미널에서 uvicorn app:app --host 0.0.0.0 --port 8000을 실행하고, 다른 터미널에서 curl로 /health와 /generate를 호출해. 예열했다면 첫 번째와 두 번째 /generate의 시간이 거의 같아야 해. 확인한 점과 다음에 더할 기능을 두 문장으로 적어.