[FastAPI] 파이썬으로 웹 구현하기 (2) | FastAPI 확장 기능 : Lifespan, APIRouter, HTTPException, Background Task

Lifesapn function

FastAPI 앱을 실행하거나 종료할 때 로직을 넣고 싶을 경우 사용합니다. 예를 들어, 앱이 시작할 때 머신런이 모델을 로드하고 앱을 종료하면 db 연결을 정리하면서 각각 상황마다 출력하도록 구현할 수 있습니다. 아래 코드는 예제를 구현한 것으로 기본적인 FastAPI 구조는 아래 글을 참고해보시기 바랍니다.

 

2024.12.12 - [Python/etc] - [Fast API] 파이썬으로 웹 구현하기 (1) | Parameter, Form, File, Request & Response Body

 

from contextlib import asynccontextmanager


def fake_answer_to_everything_ml_model(x: float):
    return x * 42

ml_models = {}

@asynccontextmanager
async def lifespan(app: FastAPI):
    print("Start Up Event")
    ml_models['answer_to_everything'] = fake_answer_to_everything_ml_model

    yield # 앱이 기동되는 시점

    print("Shutdown Event!")
    ml_models.clear()


app = FastAPI(lifespan=lifespan)


@app.get("/predict")
def predict(x: float):
    result = ml_models['answer_to_everything'](x)
    return {"result" : result}

 

아래 이미지와 같이 잘 잘동하는 것처럼 보입니다. 

 

API Router

앱이 커짐에 따라 get, post가 많아질텐데 이를 관리하기 위해 API Router를 활용할 수 있습니다. @app.get, @app.post을 사용하지 않고 별도의 router 파일을 따로 설정하고 app에 가져와서 사용하는 방식입니다. 다음 예제는 user라는 라우터를 별도로 만들어 app에 연결하는 코드입니다. 우선은 user.py로 만드는 코드입니다.

from fastapi import APIRouter


user_router = APIRouter(prefix="/users")


@user_router.get("/", tags=['users'])
def read_users():
    return [{'username':'Rick'}, {"username": "Morty"}]


@user_router.get("/me", tags=['users'])
def read_user_me():
    return {"username": "fakecurrentuser"}


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

 

다음으로 저희가 만든 앱에 해당 라우터를 불러옵니다. 같은 디렉토리에 user.py가 있다고 가정하고 import해온 후에 이를 FastAPI 앱에 include_router를 통해 추가해줍니다. 역서 참고로 include_router는 스크립트가 실행되는 __name__ == '__main__'에 들어가게 되면 app에 라우터가 포함되지 않은 상태로 정의되기 때문에 제대로 처리하지 못합니다.

from fastapi import FastAPI
import uvicorn
import user

app = FastAPI()
app.include_router(user.user_router)


if __name__ == '__main__':
    uvicorn.run(app, host="0.0.0.0", port=8000)

 

아래와 같이 잘 작동하고 있는 것을 확인할 수 있습니다.

Error Handler

에러가 발생한 경우 관련 기록을 남기고, 메시지를 클라이언트에 보내는 것은 애플리케이션을 운영하는 관점에서 필요한 일입니다. FastAPI에서는 이러한 에러 처리를 위해 HTTP Exception을 통해 가능하도록 하고 있습니다. 아래는 그 예시 코드입니다. 

 

from fastapi import HTTPException


@app.get("/v/{item_id}")
async def find_by_id(item_id: int):
    try :
        item = items[items_id]
    except KeyError:
        raise HTTPException(status_code=404, 
                            detail=f"아이템을 찾을 수 없습니다. id : {item_id}")
    return item

 

아래 이미지를 보면 제대로 입력된 경우(왼쪽) 원하는 값을 반환하지만, 사전에 정의되지 않은 오류 발생의 경우(오른쪽) 에러 메시지를 띄우는 것을 확인할 수 있습니다.

Background Task

여러 작업 중 시간이 오래 소요되는 것들에 대해서는 백그라운드로 실행할 수 있습니다. 다음은 주어진 대기 시간(wait_time) 동안 작업을 수행하고 생성한 작업 결과를 저장해 필요할 때 조회하는 코드입니다. 

from time import sleep
from fastapi import BackgroundTasks
from uuid import UUID, uuid4
from pydantic import BaseModel, Field


# 입력 데이터 모델 
class TaskInput(BaseModel):
    id_: UUID = Field(default_factory=uuid4)
    wait_time: int

# 대기 및 작업 저장
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) # 비동기 작업에 대해서 status code 202를 반환
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 "ok"

# 작업 결과 조회
@app.get("/task/{task_id}")
def get_task_result(task_id: UUID):
    try:
        return task_repo[task_id]
    except KeyError:
        return None

참고자료

[1] 변성윤. "[Product Serving] Fast API (2)". boostcamp AI Tech. 

[2] https://fastapi.tiangolo.com/advanced/events/

[3] https://fastapi.tiangolo.com/reference/apirouter/

[4] https://fastapi.tiangolo.com/tutorial/handling-errors/

[5] https://fastapi.tiangolo.com/tutorial/background-tasks/