FastAPI

J. Hwang·2024년 12월 16일

FastAPI는 파이썬 기반의 웹 프레임워크로, 빠르고 효율적으로 API 개발을 할 수 있게 해주는 도구이다.
클라이언트의 요청을 받아 즉각적으로 응답을 제시하는 online serving을 위해 API 웹을 FastAPI로 만들어보자.

프로젝트 구조

 ┣ app
 ┃  ┗ __init__.py    
 ┃  ┗ main.py         # FastAPI의 애플리케이션과 Router 설정
 ┃  ┗ config.py       # 데이터베이스, 모델 경로, 실행 환경을 정의 (configuration)
 ┃  ┗ api.py          # 모델 기능 구현
 ┃  ┗ schemas.py      # 요청/응답을 주고 받는 형태 정의
 ┃  ┗ database.py     # 데이터를 저장
 ┃  ┗ model.py        # ML model에 대한 클래스와 함수 정의

간단한 ML 모델을 로드해서 데이터를 입력하면 모델이 예측하여 그 결과를 반환해주는 간단한 API 웹을 만드는 프로젝트이다.

혹시 프로젝트 구조를 어떻게 구성해야 할지 잘 모르겠다면 아래 페이지들을 참고하자.

설치

$ pip install fastapi

config.py

from pydantic import Field
from pydantic_settings import BaesSettings

class Config(BaseSettings):
    db_url:str = Field(default="sqlite:///./db.sqlite3", env="DB_URL")
    model_path:str = Field(default="model.joblib", env="MODEL_PATH")
    app_env:str = Field(default="local", env="APP_ENV")
    
config = Config()

데이터베이스, 모델 경로, 실행 환경을 정의하는 configuration을 설정하는 파일이다. pydantic에 대해 더 알고 싶다면 여기 참조

dataset.py

import datetime
from sqlmodel import SQLModel, Field, create_engine
from config import config

class PredictionResult(SQLModel, table=True):
    id:int = Field(default=None, primary_key=True)
    result:int
    created_at:str = Field(default_factory=datetime.datetime.now)     # default_factory : default를 설정. 동적으로 값을 지정.
    
engine = create_engine(config.db_url)

모델이 예측한 결과를 테이블로 저장하기 위한 파일이다.

model.py

def load_model(model_path: str):
    import joblib

    return joblib.load(model_path)

모델과 관련된 함수들을 저장한 파일이다. 이외에도 필요에 따라 get_model, train, predict, save_model 등의 함수를 정의하여 사용하면 된다.

schemas.py

from pydantic import BaseModel

class PredictionRequest(BaseModel):
    features: list     # input 데이터를 상황에 맞게 작성
    
class PredictionResponse(BaseModel):
    id: int
    result: int

schema란 네트워크를 통해 데이터를 주고 받을 때 어떤 형태로 주고 받을지 정의하는 개념이다. 주로 예측 요청과 예측 응답을 주고 받을 형태를 정의한다.

api.py

from fastapi import APIRouter, HTTPException, status
from schemas import PredictionRequest, PredictionResponse
from model import get_model
from database import PredictionResult, engine
from sqlmodel import Session, select

router = APIRouter()

@router.post('/predict')
def predict(request: PredictionRequest) -> PredictionResponse:
    # 모델 load
    model = get_model()
    
    # 에측 : 코드 상황에 따라 구현
    prediction = int(model.predict([request.features])[0])
    
    # 예측한 결과를 DB에 저장
    # 데이터 베이스 객체를 생성. 그 때 prediction을 사용
    prediction_result = PredictionResult(result=prediction)
    with Session(engine) as session:
        session.add(prediction_result)
        session.commit()
        session.refresh(prediction_result)
        
    return PredictionResponse(id=prediction_result.id, result=prediction)

@router.get("/predict/{id}")        
def get_prediction(id: int) -> PredictionResponse:
    with Session(engine) as session:
        prediction_result = session.get(PredictionResult, id)
        if not prediction_result:
            raise HTTPException(
                detail="Not Found", status_code = status.HTTP_404_NOT_FOUND
            )
        return PredictionResponse(
            id=prediction_result.id, result=prediction_result.result
        )

main.py가 기능들을 "API 웹으로 구현"하는 파일이라면, api.py는 주로 구현할 "기능"을 정의하는 파일이다. 이 api.py 파일에서는 1) 모델을 불러와 예측을 수행하는 기능, 2) id를 입력하면 id값에 해당하는 저장된 예측 결과 값을 반환하는 기능을 구현했다.

main.py

from fastapi import FastAPI

app = FastAPI()

@app.get("/")
def root():
    return "Hello World!"
    
if __name__ == "__main__":
    import uvicorn
    
    uvicorn.run("main:app", host="0.0.0.0", port=8000, reload=True)

위는 아주 간단한 웹 API를 구현하는 코드이다. 이 코드를 만들어 저장한 뒤 웹 브라우저에서 localhost:8000으로 접속하면 "Hello World!"를 출력하는 웹 페이지가 만들어진 것을 확인할 수 있다.

이제 모델을 로드할 수 있도록 코드를 추가하자.

from contextlib import asynccontextmanager
from loguru import logger
from sqlmodel import SQLModel

from config import config
from database import engine
from model import load_model
from api import router

@asynccontextmanager
async def lifespan(app: FastAPI):
    # 데이터 베이스 테이블 생성
    logger.info("Creating database table")
    SQLModel.metadata.create_all(engine)
    
    # 모델 로드 
    logger.info("Loading model")
    load_model(config.model_path)     # model.py에 존재 
    yield
    
app = FastAPI(lifespan=lifespan)
app.include_router(router)

위 코드를 포함해 main.py를 실행시키면 데이터 베이스 테이블이 생성되면서 FastAPI 디렉토리에 db.sqlite3 파일이 생성됨을 확인할 수 있다.

위의 lifespan 함수는 FastAPI 앱의 시작과 종료 단계에서 실행할 특정 작업을 정의하는데 사용한다. yield를 기준으로 yield 이전 라인은 앱을 시작하기 전에 실행할 작업이고 yield 이후 라인은 앱을 종료하기 전에 실행할 작업이다. 위의 예시에서는 앱이 실행되기 전 데이터 베이스 테이블이 생성되고 모델이 로드된다. 앱을 종료하기 전 작업은 정의되지 않은 상태이다.

동작 확인하기

데이터를 입력하면 모델이 예측하여 그 결과를 반환해주는 API 웹이 완성되었으니 이제 동작을 확인해보자.
우선 CLI에서 예측을 요청해보자. (명령어가 이해되지 않는다면 여기 참조)

$ curl -X POST "http://0.0.0.0:8000/predict" -H "Content-Type: application/json" -d "{'features': [5.1, 3.5, 1.4, 0.2]}"

이는 방금 만든 예측 서버에 [5.1, 3.5, 1.4, 0.2]라는 데이터를 입력해준 것이다. 웹이 정상적으로 만들어졌다면 {"id":0, "result":0}과 같이 예측 결과를 반환해준다.

(ps. 여기서는 아주 간단한 구현을 위해 iris 데이터를 sklearn의 모델로 예측하도록 구성하였다.)

그리고 api.py에서 볼 수 있듯이 예측한 결과를 불러올 수도 있게 구현했기 때문에 이 또한 CLI에서 확인해보자.

$ curl -X GET "http://0.0.0.0:8000/predict/0"

이는 예측 결과가 저장된 데이터 베이스에서 0번 예측 결과를 가져오라고 요청한 것이다. 웹이 정상적으로 작동한다면 {"id":0, "result":0}과 같이 저장된 예측 결과를 잘 불러올 것이다.

docs 확인하기

구현한 API 웹에서 docs 페이지(localhost:8000/docs)에 접속하면 지금까지 구현한 기능들을 살펴볼 수 있다.

기타 기능 구현하기

FastAPI에서는 HTTP 메서드를 데코레이터로 표시한다. 이를 이용해 다양한 기능이 구현되도록 할 수 있다.

Path parameter

@app.get("/users/{user_id}")
def get_user(user_id):
    return {"user_id": user_id}

path parameter를 이용해 user id를 조회하는 기능을 구현한 함수이다.
main.py에 포함시켜 사용한다.

Query parameter

items_db = [{"item_name": "Apple"}, {"item_name": "Banana"}, {"item_name": "Cake"}]

@app.get("/items/")
def read_item(skip: int = 0, limit: int = 10):
    return items_db[skip: skip + limit]

query parameter를 이용해 item을 조회/검색하는 기능을 구현한 함수이다.
localhost:8000/items/로 접근하여 전체 item을 살필 수도 있고, localhost:8000/items?item_name=Apple과 같이 검색할 수도 있다.

Form

python의 input 함수처럼 입력 형태로 데이터를 받고 싶은 경우가 있다. 이를 Form이라고 하는데, 이 기능을 사용하려면 python-multipart를 설치해야 한다.

pip install python-multipart

from fastapi import FastAPI, Form, Request

@app.post("/login/")
def login(username: str = Form(...), password: str = Form(...)):
    return {"username": username}

python-multipart를 설치하고 위와 같은 함수를 작성하여 main.py에 포함시키면 로그인하는 기능을 구현할 수 있다. 사용자로부터 id와 password를 입력받아 로그인 요청을 처리하는 것이다.

프론트엔드 구현하기

Jinja2를 이용하면 FastAPI로 만든 백엔드 웹을 프론트엔드로 구현할 수 있다.

pip install Jinja2

from fastapi.templating import Jinja2Templates

templates = Jinja2Templates(directory='./')

@app.get("/login/")
def get_login_form(request: Request):
    return templates.TemplateResponse('login_form.html', context={'request': request})

위 코드를 main.py에 작성하면 직전에 만든 login 기능을 프론트엔드로 구성할 수 있다.
이 때 로그인 페이지의 프론트엔드 구성을 정의하는 login_form.html이라는 파일을 같은 디렉토리에 저장해둬야 한다.

그러면 위와 같이 아이디와 비밀번호를 입력하는 창과 제출 버튼이 나타나는 것을 확인할 수 있다.

파일 업로드

API 웹에 파일을 업로드하도록 구현할 수 있다. Form을 구현할 때처럼 python-multipart 설치가 필요하다.

from typing import List

from fastapi import File, UploadFile
from fastapi.responses import HTMLResponse

@app.post("/files/")
def create_files(files: List[bytes] = File(...)):
    return {"file_sizes": [len(file) for file in files]}


@app.post("/uploadfiles/")
def create_upload_files(files: List[UploadFile] = File(...)):
    return {"filenames": [file.filename for file in files]}


@app.get("/")
def main():
    content = """
<body>
<form action="/files/" enctype="multipart/form-data" method="post">
<input name="files" type="file" multiple>
<input type="submit">
</form>
<form action="/uploadfiles/" enctype="multipart/form-data" method="post">
<input name="files" type="file" multiple>
<input type="submit">
</form>
</body>
    """
    return HTMLResponse(content=content)

위 코드를 작성하여 main.py에 포함시키면 파일을 업로드하는 기능을 구현할 수 있다.

API router

프로젝트가 크고 복잡해지면 API endpoint가 점점 많아져서 @app.get, @app.post와 같은 코드를 하나의 모듈에서 관리하기 어려워진다. 따라서 많은 API endpoint를 효율적으로 관리하기 위해 API router를 사용한다. router.py 라는 파일을 새로 생성하여 안에 user라는 router와 order라는 router를 정의하는 코드를 작성하여 저장한다.

from fastapi import APIRouter

user_router = APIRouter(prefix="/users")
order_router = APIRouter(prefix="/orders")

@user_router.get("/{username}", tags=["users"])
def read_user(username: str):
    return {"username": username}
    
@order_router.get("/{order_id}", tags=["orders"])
def read_order_id(order_id: str):
    return {"order_id": order_id}

(프로젝트가 더 크고 복잡해지는 경우에는 router.py에 user와 order를 같이 포함시키기 보다는 user.py, order.py로 분리하여 router를 작성한다.)

그리고 main.py에는 아래와 같이 router를 import 해서 등록해주면 된다.

from router import user_router, order_router

app.include_router(user_router)
app.include_router(order_router)

Background task 정의하기

오래 걸리는 작업을 background task로 실행하면 기다리지 않고 바로 응답을 줄 수 있다. 이런 background task를 정의하고 싶다면 async를 이용하면 된다.

from uuid import UUID, uuid4

class TaskInput(BaseModel):
    id_: UUID = Field(default_factory=uuid4)
    wait_time: int


task_repo = {}


def cpu_bound_task(id_: UUID, wait_time: int):
    sleep(wait_time)
    result = f"task done after {wait_time}"
    task_repo[id_] = result


@app.post("/task", status_code=202)   # 비동기 작업이 등록됐을 때 보통 HTTP Response 202 (Accepted)를 리턴
async def create_task_in_background(task_input: TaskInput, background_tasks: BackgroundTasks):
    background_tasks.add_task(cpu_bound_task, id_=task_input.id_, wait_time=task_input.wait_time)
    return task_input.id_

# 작업 결과물을 저장해뒀다가 get 요청을 통해 완료 확인
@app.get("/task/{task_id}")
def get_task_result(task_id: UUID):
    try:
        return task_repo[task_id]
    except KeyError:
        return None