우리가 FastAPI 와 같은 Application 을 만들다보면 요청 전/후로 리소스를 정리하거나 미리 무거운 리소스를 로드하는 등의 작업이 필요하게 된다. FastAPI 에서는 이를 asynccontextmanager 를 이용한 lifespan 으로 지원해주는데 오늘은 이를 알아보도록 하자.

ContextManager

asynccontextmanager 에 관해서 알아보기 전에 우리는 ContextManager 에 대한 개념이 필요하다. ContextManager 는 with 문과 함께 실행시간에 의존하는 컨텍스트를 생성하게 된다.

따라서 메소드도 __enter__() 와 __exit__()(이와 같이 파이썬에서 underbar 두개를 붙이면 던더메소드¹ 라고 한다) 을 가지고 있다. 즉, 컨텍스트에 진입할때 __enter__() 가 호출되고 컨텍스트와 관련된 내용을 반환하고, 처리를 마친뒤에는 exit() 함수가 호출되게 된다. 아래 코드 예시를 보면 조금 더 쉬울 것 이다.

class ManagedResource:
    def __init__(self, *args, **kwds):
        self.args = args
        self.kwds = kwds
        self.resource = None

    def __enter__(self):
        # 자원을 얻는 코드, 예를 들어:
        self.resource = acquire_resource(*self.args, **self.kwds)
        return self.resource

    def __exit__(self, exc_type, exc_val, exc_tb):
        # 자원을 해제하는 코드, 예를 들어:
        release_resource(self.resource)

보면 __enter__ 에서 리소스를 취득한 뒤에 리소스(컨텍스트와 연관된)를 반환 받고, 받는 쪽에서 호출이 끝난 뒤 __exit__ 에서 realeas_resource(resource) 함수가 호출된다. 이러한 contextmanager 를 조금 더 쉽게 만들기 위한 팩터리 메서드² 를 호출하는 데코레이터 방식을 이용하면 조금 더 간결하게 구현이 가능하다. 아래 예시를 한번 보자.

from contextlib import contextmanager

@contextmanager
def managed_resource(*args, **kwds):
    # 자원을 얻는 코드, 예를 들어:
    resource = acquire_resource(*args, **kwds)
    try:
        yield resource
    finally:
        # 자원을 해제하는 코드, 예를 들어:
        release_resource(resource)

@contextmanager 와 같은 데코레이터 방식은 이러한 컨텍스트매니저를 편리하게 만들 수 있도록 도와주는 팩터리 메서드이다. 조금 더 직관적으로 확인하기 위해 코드로 한번 확인해보자.

import asyncio
from contextlib import asynccontextmanager

@asynccontextmanager
async def hello():
  print("before executing function")
  yield "hello"
  print("after executing function")

async def main():
  async with hello() as data:
    print(data)

asyncio.run(main())

이 함수의 실행결과는 어떨까? 아래에서 한번 확인해보자.

before executing function
hello
after executing function

위와 같이 실행되는 이유는 yield 의 윗 부분이 __enter__ 의 역할을 하고 yield 값을 반환하고 난 뒤, 해당 값을 받아 외부에서 사용하고 다시 문맥이 종료되는 시점(with 문이 끝나는 시점)에 yield 하위 코드가 실행된다. 코드를 보면 어렵지 않고 코루틴에 익숙하다면 yield 또한 익숙할 것이다. 그렇다면 FastAPI 는 이 asynccontextmanager 를 어떻게 이용해서 구현하고 있을까? 예상이 가지 않는가?

lifespan

@asynccontextmanager
async def lifespan(app: FastAPI):
    # Load the ML model
    ml_models["answer_to_everything"] = fake_answer_to_everything_ml_model
    yield
    # Clean up the ML models and release the resources
    ml_models.clear()


app = FastAPI(lifespan=lifespan)

예시를 보면 알겠지만 윗 내용을 이해했다면 아주 쉬울 것이다. yield 윗 부분의 Model 을 로딩하는 부분이 FastAPI application 이 시작되기전에 실행되며, yield 하위 부분은 FastAPI 실행이 끝나면 모델을 정리하게 된다. 즉, 간단하게 표시해보면 아래와 같은 코드임을 생각해 볼 수 있다.

async with lifespan():
  fastapi.run()