Bài 5: Load model 1 lần khi startup (lifespan event)

Sau bài này bạn sẽ:

• ✅ Hiểu tại sao không được load model bên trong handler mỗi request
• ✅ Sử dụng lifespan context manager (FastAPI 0.93+) để load model 1 lần khi server start
• ✅ Chia sẻ model giữa các request qua global dict hoặc app.state
• ✅ Inject model vào endpoint sạch sẽ bằng Depends()
• ✅ Hiểu tradeoff RAM khi chạy nhiều worker
• ✅ Thêm warm-up inference để tránh latency cao cho request đầu tiên

Model sklearn nhỏ (~10 MB) load mất vài chục ms. Model PyTorch trung bình (~200 MB) mất 1–3 giây. Model Transformers lớn (1–7 GB) mất 10–60 giây. Nếu mỗi request gọi lại load_model(), mỗi client phải chờ thêm ngần ấy thời gian — hoàn toàn không khả thi.

Ví dụ code sai:

# KHÔNG làm thế này
from fastapi import FastAPI
import joblib

app = FastAPI()

@app.post("/predict")
def predict(features: list[float]):
    clf = joblib.load("model.pkl")   # load lại mỗi request → chậm
    return {"pred": clf.predict([features]).tolist()}

Với 100 request/giây, mỗi request mất 500 ms chỉ để load model, server chịu thêm 50 lần I/O disk và RAM allocation không cần thiết mỗi giây.

Giải pháp: load model 1 lần khi server khởi động, giữ trong bộ nhớ, chia sẻ object đó cho mọi request.

Trước FastAPI 0.93 (phát hành 2023-03), cách phổ biến là dùng decorator @app.on_event:

# Pattern CŨ — deprecated từ FastAPI 0.93
@app.on_event("startup")
async def load_model():
    app.state.clf = joblib.load("model.pkl")

@app.on_event("shutdown")
async def cleanup():
    del app.state.clf

Code trên vẫn chạy được nhưng FastAPI sẽ in cảnh báo DeprecationWarning mỗi lần khởi động. Bài này chỉ dạy pattern mới là lifespan context manager.

Từ FastAPI 0.93+ (Starlette 0.26+), cách chuẩn là dùng @asynccontextmanager:

from contextlib import asynccontextmanager
from fastapi import FastAPI
import joblib

# Dict toàn cục — đủ dùng cho hầu hết use case
ml_models: dict = {}

@asynccontextmanager
async def lifespan(app: FastAPI):
    # --- Startup ---
    # Chạy TRƯỚC khi server accept request đầu tiên
    ml_models["clf"] = joblib.load("model.pkl")
    print("Model loaded.")

    yield  # Server nhận request ở đây

    # --- Shutdown ---
    # Chạy SAU khi server dừng (SIGTERM, Ctrl+C)
    ml_models.clear()
    print("Model released.")

app = FastAPI(lifespan=lifespan)

@app.post("/predict")
def predict(features: list[float]):
    clf = ml_models["clf"]
    return {"pred": clf.predict([features]).tolist()}

Cơ chế hoạt động:

• Code trước yield: chạy khi server start, trước khi accept bất kỳ request nào. Nếu đoạn này raise exception, server không start.
• yield: điểm bàn giao — server bắt đầu nhận request và chạy cho đến khi có tín hiệu dừng.
• Code sau yield: chạy khi server nhận SIGTERM hoặc Ctrl+C. Phù hợp để giải phóng resource (đóng kết nối DB, xóa cache tạm).

Lifespan là async context manager, nên bạn có thể dùng await bên trong — hữu ích khi load model từ remote (S3, Hugging Face Hub):

@asynccontextmanager
async def lifespan(app: FastAPI):
    # Load từ Hugging Face Hub bất đồng bộ (giả sử dùng httpx)
    model_bytes = await download_model_async("hf://org/model")
    ml_models["model"] = deserialize(model_bytes)
    yield
    ml_models.clear()

Thay vì dict global, FastAPI cung cấp app.state — một namespace gắn vào object app, truy cập được từ request.app.state trong bất kỳ endpoint nào.

from contextlib import asynccontextmanager
from fastapi import FastAPI, Request
from pydantic import BaseModel
import joblib

class PredictRequest(BaseModel):
    features: list[float]

@asynccontextmanager
async def lifespan(app: FastAPI):
    app.state.clf = joblib.load("model.pkl")
    yield
    del app.state.clf   # giải phóng bộ nhớ

app = FastAPI(lifespan=lifespan)

@app.post("/predict")
def predict(req: Request, body: PredictRequest):
    clf = req.app.state.clf
    prediction = clf.predict([body.features])
    return {"pred": prediction.tolist()}

So sánh hai cách:

• Global dict: đơn giản, đủ dùng. Nhược điểm: dict global có thể bị code khác trong module vô tình ghi đè.
• app.state: namespace riêng của app, ít rủi ro xung đột tên hơn. Phù hợp khi có nhiều model hoặc nhiều resource (DB pool, config).

Lưu ý: app.state là starlette.datastructures.State — truy cập attribute tùy ý, không cần khai báo trước.

Truy cập req.app.state.clf trực tiếp trong mỗi handler gây lặp code. Cách sạch hơn là bọc vào một dependency function rồi dùng Depends():

from contextlib import asynccontextmanager
from fastapi import FastAPI, Request, Depends
from pydantic import BaseModel
import joblib
from sklearn.base import ClassifierMixin

class PredictRequest(BaseModel):
    features: list[float]

@asynccontextmanager
async def lifespan(app: FastAPI):
    app.state.clf = joblib.load("model.pkl")
    yield
    del app.state.clf

app = FastAPI(lifespan=lifespan)

# Dependency function — trả về classifier từ app.state
def get_clf(request: Request) -> ClassifierMixin:
    return request.app.state.clf

@app.post("/predict")
def predict(body: PredictRequest, clf: ClassifierMixin = Depends(get_clf)):
    prediction = clf.predict([body.features])
    return {"pred": prediction.tolist()}

Lợi ích của pattern này:

• Handler không biết model đến từ đâu — dễ mock trong test (app.dependency_overrides).
• Thêm model mới chỉ cần thêm dependency function mới, không sửa handler cũ.
• Type hint rõ ràng, IDE autocomplete hoạt động tốt hơn.

Ví dụ thực tế với PyTorch

from contextlib import asynccontextmanager
from fastapi import FastAPI, Request, Depends
from pydantic import BaseModel
import torch
import torch.nn as nn

# Giả sử đã định nghĩa SimpleClassifier
class SimpleClassifier(nn.Module):
    def __init__(self):
        super().__init__()
        self.fc = nn.Linear(4, 3)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        return self.fc(x)

DEVICE = "cuda" if torch.cuda.is_available() else "cpu"

@asynccontextmanager
async def lifespan(app: FastAPI):
    model = SimpleClassifier()
    # Load weights từ checkpoint
    state_dict = torch.load("classifier.pt", map_location=DEVICE)
    model.load_state_dict(state_dict)
    model.to(DEVICE)
    model.eval()   # tắt dropout, batchnorm ở chế độ inference
    app.state.model = model
    app.state.device = DEVICE
    yield
    del app.state.model

app = FastAPI(lifespan=lifespan)

class InferRequest(BaseModel):
    features: list[float]   # ví dụ: 4 features

def get_model(request: Request) -> nn.Module:
    return request.app.state.model

def get_device(request: Request) -> str:
    return request.app.state.device

@app.post("/classify")
@torch.no_grad()
def classify(
    body: InferRequest,
    model: nn.Module = Depends(get_model),
    device: str = Depends(get_device),
):
    x = torch.tensor([body.features], dtype=torch.float32).to(device)
    logits = model(x)
    pred_class = logits.argmax(dim=-1).item()
    return {"class": pred_class}

Lưu ý quan trọng với PyTorch:

• Gọi model.eval() sau khi load — tắt Dropout và BatchNorm ở chế độ train.
• Dùng @torch.no_grad() hoặc bọc trong with torch.no_grad(): để tắt gradient computation khi inference — tiết kiệm bộ nhớ và nhanh hơn.
• torch.load(..., map_location=DEVICE) tránh lỗi khi model được train trên GPU nhưng server inference dùng CPU.

Khi chạy với nhiều worker:

uvicorn main:app --workers 4

Uvicorn fork ra 4 process Python độc lập. Mỗi process là 1 Python interpreter riêng, có bộ nhớ riêng. lifespan chạy trong mỗi worker — tức model được load 4 lần vào 4 vùng RAM khác nhau.

Hệ quả:

• Model 2 GB + 4 workers = ~8 GB RAM cho model (chưa kể overhead).
• Throughput tăng (4 workers xử lý song song) nhưng RAM tăng tuyến tính theo số worker.

Các hướng giảm RAM:

• 1 worker + nhiều thread: phù hợp với sklearn và model nhẹ (GIL vẫn là giới hạn với Python thuần).
• Quantization: giảm model từ FP32 xuống INT8/INT4 — xem bài 55.
• Inference server riêng: TorchServe, Triton Inference Server, hay vLLM cho LLM — các tool này quản lý model tập trung, nhiều worker API chỉ gọi qua gRPC/HTTP. Chủ đề này nằm ngoài scope module FastAPI.

Quy tắc đơn giản: nếu RAM của server cho phép, dùng nhiều worker để tăng throughput. Nếu RAM là bottleneck, giữ ít worker hoặc dùng inference server chuyên dụng.

Lần inference đầu tiên thường chậm hơn các lần sau vì:

• CUDA kernel compile (JIT): PyTorch compile kernel CUDA lần đầu khi gặp shape tensor mới.
• Memory allocation: GPU cấp phát vùng nhớ cho activations, gradients (nếu dùng).
• CPU cache cold start: với sklearn trên CPU, cache processor chưa có data model.

Để tránh user đầu tiên chịu latency cao, chạy 1 lần inference dummy trong lifespan sau khi load model:

import torch

@asynccontextmanager
async def lifespan(app: FastAPI):
    model = load_pytorch_model("weights.pt")
    model.to(DEVICE)
    model.eval()

    # Warm-up: inference với input dummy có cùng shape thực tế
    # Ví dụ model nhận ảnh 224x224 RGB, batch size 1
    dummy_input = torch.zeros(1, 3, 224, 224).to(DEVICE)
    with torch.no_grad():
        _ = model(dummy_input)

    app.state.model = model
    yield
    del app.state.model

Warm-up nên dùng cùng shape với input thực tế. Nếu model nhận shape động, chạy warm-up với shape phổ biến nhất.

Với Transformers (Hugging Face), warm-up tương tự:

from transformers import pipeline

@asynccontextmanager
async def lifespan(app: FastAPI):
    pipe = pipeline("text-classification", model="distilbert-base-uncased-finetuned-sst-2-english")

    # Warm-up
    _ = pipe("warm up")

    app.state.pipe = pipe
    yield
    del app.state.pipe

1. Gán lifespan sau khi tạo app

# SAI — không work
app = FastAPI()
app.lifespan = lifespan   # attribute này không tồn tại

# ĐÚNG — truyền vào constructor
app = FastAPI(lifespan=lifespan)

FastAPI đọc lifespan khi khởi tạo object và pass xuống Starlette. Gán sau không có tác dụng và không báo lỗi.

2. Startup quá chậm gây health check fail

Load balancer (AWS ALB, GCP GLB, Kubernetes liveness probe) thường gửi health check request ngay sau khi container start. Nếu lifespan startup mất quá lâu (ví dụ download model 5 GB), server chưa kịp nhận request và load balancer có thể đánh dấu instance là unhealthy.

Giải pháp:

• Bake model vào Docker image thay vì download khi start.
• Tăng initialDelaySeconds (Kubernetes) hoặc health check grace period của load balancer.
• Dùng readinessProbe riêng — chỉ chuyển traffic vào sau khi model sẵn sàng.

3. Quên cleanup gây memory leak khi reload

Uvicorn có chế độ --reload cho dev. Mỗi lần reload, lifespan chạy lại. Nếu code sau yield không giải phóng model, mỗi lần reload tích lũy thêm một bản model trong RAM:

# Dễ gây leak khi --reload
@asynccontextmanager
async def lifespan(app: FastAPI):
    app.state.model = load_heavy_model()
    yield
    # Quên del app.state.model → model cũ còn đó sau reload

# Đúng
@asynccontextmanager
async def lifespan(app: FastAPI):
    app.state.model = load_heavy_model()
    yield
    del app.state.model   # giải phóng reference

4. Model không thread-safe

Một số model hoặc thư viện inference không thread-safe — gọi đồng thời từ nhiều thread có thể gây race condition hoặc crash. Với FastAPI sync endpoint (def predict), FastAPI chạy trong thread pool. Với async endpoint (async def predict), FastAPI chạy trong event loop của process đó.

Kiểm tra tài liệu của library bạn dùng. PyTorch inference (forward pass) thread-safe từ PyTorch 1.9+. Scikit-learn predict thread-safe cho hầu hết estimator.

5. Dùng lifespan với router phụ (APIRouter)

APIRouter không có lifespan riêng — lifespan chỉ đặt ở FastAPI(...) chính. Các router được include vào app chính sẽ tự dùng chung lifespan của app đó.

Bảng so sánh các cách chia sẻ model:

Checklist triển khai:

• ✅ Dùng FastAPI(lifespan=lifespan), không dùng @app.on_event
• ✅ Load model trước yield, cleanup sau yield
• ✅ Gọi model.eval() và bọc inference trong torch.no_grad() với PyTorch
• ✅ Chạy warm-up inference trong lifespan nếu cần latency ổn định
• ✅ Tính toán RAM = model_size × số_worker trước khi chọn cấu hình server
• ✅ Không quên del hoặc .clear() trong phần shutdown

Bài 6: Endpoint inference — nhận data, trả kết quả prediction