Backend.AI에서 사용되는 GraphQL batching과 Pagination

Backend.AI에서는 GraphQL을 사용하고 있습니다. 이번 포스트에서는 크게 세 가지 주제에 대해서 이야기 해보려고 합니다. 첫 번째로 GraphQL에 대한 간단한 설명과 GraphQL을 어떻게 사용하고 있는지에 대해서, 두 번째로 GraphQL을 사용하면서 어떤 기술적 문제에 봉착했고 어떻게 해결했는지, 마지막으로 Pagination에 대해서 설명하는 시간을 가지려 합니다.

GraphQL 개괄

GraphQL은 페이스북에서 만든 쿼리 언어입니다. 페이스북 화면에는 굉장히 다양한 정보들이 있습니다. 개인의 정보, 타임라인, 알림 메시지 등 이런 것들을 개별적으로 쿼리 하게 된다면 쿼리마다 수십 개의 connection이 만들어지고 이 모두를 JSON REST API로 가져오려면 round trip의 횟수가 매우 큽니다(over-fetching, under-fetching, REST 엔드포인트 관리 등)¹. GraphQL은 REST API가 가지는 위의 단점들을 극복하여 더 빠르고 가벼운 웹 애플리케이션을 만들 수 있게 도와줍니다.

GraphQL을 래블업에서는 어떻게 사용하고 있는가?

Backend.AI에서 GraphQL을 사용하는 이유는 앞서 언급한 round-trip 횟수 감소라는 장점 외에도 리졸버(resolver) 구현을 유연하게 할 수 있기 때문입니다. 유연하게 사용한다는 것은 두 가지 의미로 설명할 수 있습니다. 첫 번째로는 예전 버전의 오브젝트 스키마에는 있었지만 새 버전의 스키마에서는 없어진 필드들을 오브젝트 생성 이후 별도의 리졸버 메소드를 통해 새 버전의 오브젝트로부터 내용을 유추·변환(interpolate)하여 채우고, 두 번째로는 하나의 오브젝트에 속하는 필드들을 서로 상이한 데이터 소스로부터 가져와 조합하여 일관된 데이터 모델을 제공할 수 있게 해줍니다.

첫 번째에 대해서 예시로는 아래의 resolve_legacy_compute_session() 메소드처럼 새로운 버전에서 없어진 필드에 대한 리졸버를 따로 추가 정의해주면, 현재 신버전 테이블의 채워진 내용으로부터 과거 레거시를 위한 테이블 내용을 채워줄 수 있습니다. 이러한 점은 버전이 업데이트 되고 field, object type 등이 추가되더라도 하위호환을 지원하기에 용이합니다.

async def resolve_legacy_compute_session(
    executor: AsyncioExecutor,
    info: graphene.ResolveInfo,
    sess_id: str,
    *,
    domain_name: str = None,
    access_key: AccessKey = None,
    status: str = None,
) -> Optional[LegacyComputeSession]:
    graph_ctx: GraphQueryContext = info.context
    loader = graph_ctx.dataloader_manager.get_loader(
        graph_ctx,
        'LegacyComputeSession.detail',
        domain_name=domain_name,
        access_key=access_key,
        status=status
    )
    matches = await loader.load(sess_id)
    if len(matches) == 0:
        return None
    elif len(matches) == 1:
        return matches[0]
    else:
        raise TooManyKernelsFound

두 번째 예로는 쿼리를 통해 DB에서 가져온 row object들을 GraphQL 엔진이 이해할 수 있도록 아래처럼 from_row()에서 변환하게끔 설계하였는데,

@classmethod
def from_row(
    cls,
    ctx: GraphQueryContext,
    row: Row,
) -> Agent:
    mega = 2 ** 20
    return cls(
        id=row['id'],
        status=row['status'].name,
        status_changed=row['status_changed'],
        region=row['region'],
        
        ...
        
        # legacy fields
        mem_slots=BinarySize.from_str(row['available_slots']['mem']) // mega,
        cpu_slots=row['available_slots']['cpu'],
        gpu_slots=row['available_slots'].get('cuda.device', 0),
        tpu_slots=row['available_slots'].get('tpu.device', 0),

        ...
    )

오브젝트 필드에는 정의되어있지만 변환할 때 모두 넣어주지 못하는 경우 아래의 resolve_live_stat() 메소드처럼 별도의 리졸버를 둬서 DB가 아니라 Redis에서 가지고 올 수 있습니다. 따라서 실제로 GraphQL 스키마 상으로는 같은 객체에 있는 필드이지만 그 값은 DB에서 가져온 것일 수도 있고, Redis에서 가져온 것일 수도 있습니다. 정리하자면 GraphQL 상에서는 개념적 데이터 모델을 구현하고 실제 데이터 모델에 채워지는 것들은 다양한 데이터 소스(SQL 데이터베이스 혹은 Redis 등)에서 채워지게끔 구현할 수 있습니다.

async def resolve_live_stat(self, info: graphene.ResolveInfo) -> Any:
    ctx: GraphQueryContext = info.context
    rs = ctx.redis_stat
    live_stat = await redis.execute_with_retries(
            lambda: rs.get(str(self.id), encoding=None)
        )
    if live_stat is not None:
        live_stat = msgpack.unpackb(live_stat)
    return live_stat

GraphQL을 사용한 데이터베이스 쿼리의 비효율성 문제와 batching²을 통한 해결방법

예를 들어서 10개의 유저 데이터를 가져오고 싶을 때, GraphQL에서는 리졸버 내에서 SQL 쿼리를 작성하면, 10번의 쿼리를 전달하게 됩니다. 이는 규모가 작을 시에는 문제가 없겠지만, 전달하려는 쿼리의 개수가 늘어났을 때는 매우 비효율적입니다. 이를 해결하기 위해서 DataLoader라는 라이브러리를 이용합니다. DataLoader는 같은 GraphQL 객체를 동일하게 여러 번 리졸브하게 되면(예: 같은 리졸브 함수에 10개의 다른 user.id를 인자로 주는 경우) 쿼리를 개별적으로 처리하지 않습니다. 입력받은 인자(여기에서는 user.id)를 묶어서 리스트로 만든 다음 리졸버를 구현하면, batching을 통해서 처리합니다. 물론 batching의 개수도 개발자가 정하기 나름입니다. (Backend.AI에서는 128개를 묶어서 처리합니다) 따라서 리졸버 쿼리문의 WHERE절이 다음과 같이 변하게 됩니다.

batching 이전

WHERE users.id == $user_id

batching 이후

WHERE users.id IN $user_ids

따라서 DataLoader 라이브러리는 SQL 쿼리를 batching하는 효과를 냅니다. SQL 문장 파싱, 쿼리 계획 작성 등 데이터베이스가 개별 SQL 쿼리마다 해야 하는 다양한 일들을 1회로 줄일 수 있기 때문에 이 라이브러리를 잘 활용하는 것이 실제 SQL 데이터베이스와 round trip 횟수를 줄여줍니다. GraphQL을 통해서 클라이언트와 API 서버 간의 round trip 횟수를 줄여놨는데 쿼리를 batching하지 않는다면 쿼리의 개수만큼 round trip 횟수가 증가하므로 GraphQL의 장점이 상쇄된다고 할 수 있습니다. Backend.AI 같은 경우에는 Graphene이라는 Python용 GraphQL 엔진을 감싼 라이브러리를 사용하며, Graphene이 비동기와 동기 모두 지원하므로 비동기 리졸버인 경우에는 batching을 위해 aiodataloader를 사용하고 있습니다.

하나의 API 요청 내에서 동일 리졸버에 대한 DataLoader를 최대한 재활용함으로써 batching 효과를 최대화하고자 우리는 DataLoader 라이브러리에 추가적인 최적화를 하였습니다. 특히 GraphQL 쿼리를 단일 API 요청에 여러 개를 실어보내거나, nested object의 경우 동일한 종류의 객체를 서로 다른 tree에서 호출한다거나 하는 경우에도 동일한 dataloader가 일정하게 재사용되도록 해줍니다. Backend.AI에서는 최적화를 위해 DataLoaderManager라는 클래스를 만들어서 하나의 API 요청 안에서 DataLoaderManager 객체를 하나 만들고 이것을 반복해서 호출하는 경우 캐싱해서 가지고 있도록 디자인하였습니다.

class DataLoaderManager:

    cache: Dict[int, DataLoader]

    def __init__(self) -> None:
        self.cache = {}
        self.mod = sys.modules['ai.backend.manager.models']

    @staticmethod
    def _get_key(otname: str, args, kwargs) -> int:
        key = (otname, ) + args
        for item in kwargs.items():
            key += item
        return hash(key)

    def get_loader(self, context: GraphQueryContext, objtype_name: str, *args, **kwargs) -> DataLoader:
        k = self._get_key(objtype_name, args, kwargs)
        loader = self.cache.get(k)
        if loader is None:
            objtype_name, has_variant, variant_name = objtype_name.partition('.')
            objtype = getattr(self.mod, objtype_name)
            if has_variant:
                batch_load_fn = getattr(objtype, 'batch_load_' + variant_name)
            else:
                batch_load_fn = objtype.batch_load
            loader = DataLoader(
                apartial(batch_load_fn, context, *args, **kwargs),
                max_batch_size=128,
            )
            self.cache[k] = loader
        return loader

Pagination

pagination 구현 방법에 대해 GraphQL 공식 문서에 권장 방법이 서술되어 있지만³, 강제적인 표준은 아니기 때문에 구현체마다 규칙이 다릅니다. 저희는 limit과 offset을 사용하는 방식을 채택하였습니다. 객체가 10~20개면 한 번에 불러오면 되지만 수백, 수천 개가 넘어가게 되면 DB에서 가져오기에 부하가 커지기 때문에 pagination을 이용합니다. Backend.AI를 초기에 GraphQL로 구현할 때 pagination 없이 graphene.List라는 것을 사용하면서 적당한 필터 조건을 넣어줬습니다. 이것을 리졸브하는 load_all()이라는 메소드를 이용하여 SQL 쿼리를 한 번에 보냈습니다. 하지만 에이전트가 수십 개로 증가하고, 클라우드도 수천 명이 넘어감으로써 pagination을 해야 할 필요성이 생겼습니다. 이를 위해서 AgentList라는 새로운 object type을 정의하고 이는 PaginatedList인터페이스를 따릅니다.

class AgentList(graphene.ObjectType):
    class Meta:
        interfaces = (PaginatedList, )

    items = graphene.List(Agent, required=True)

class Item(graphene.Interface):
    id = graphene.ID()

class PaginatedList(graphene.Interface):
    items = graphene.List(Item, required=True)
    total_count = graphene.Int(required=True)

AgentList라는 object type을 만들고 리졸버 또한 구현하였습니다. resolve_agent_list()는 공통 인자로 limit, offset, filter, order을 받습니다. limit과 offset은 SQL 문법에서 쿼리 결과의 범위를 제한하는 것과 동일한 의미이고 filter와 order는 임의 필드들에 대한 조건식을 조합할 수 있는 표현식입니다. total_count는 load_count()함수를 통해서 조건에 맞는 총 에이전트의 개수를 의미하고, agent_list는 load_slice()함수를 통해 limit과 offset조건에 해당하는 에이전트들을 가져와서 채워진 리스트를 의미합니다.

async def resolve_agent_list(
    executor: AsyncioExecutor,
    info: graphene.ResolveInfo,
    limit: int,
    offset: int,
    *,
    filter: str = None,
    order: str = None,
    scaling_group: str = None,
    status: str = None,
) -> AgentList:
    total_count = await Agent.load_count(
        info.context,
        scaling_group=scaling_group,
        raw_status=status,
        filter=filter,
    )
    agent_list = await Agent.load_slice(
        info.context, limit, offset,
        scaling_group=scaling_group,
        raw_status=status,
        filter=filter,
        order=order,
    )
    return AgentList(agent_list, total_count)

정리

이번 포스트에서는

• GraphQL 개괄
• GraphQL을 Backend.AI에서 어떻게 사용하고 있는지
• GraphQL을 사용할 때의 문제점 및 batching을 통한 해결방안
• Pagination에 대한 필요성 및 구현법

위와 같은 주제들로 이야기해 보았습니다. 이 글이 GraphQL을 도입하고자 할 때 참고할 활용 및 최적화 사례로 여러분들에게 도움이 되기를 바랍니다.

note

래블업에서는 이러한 엔지니어링을 함께 할 인재를 모시고 있습니다. 👉 지원하기

---

1. https://www.howtographql.com/basics/1-graphql-is-the-better-rest/↩
2. batching이란 데이터를 요청이 들어온 즉시 처리하는 것이 아니라, 요청을 일괄적으로 모아서 처리하는 작업을 의미합니다.↩
3. https://graphql.org/learn/pagination/↩