카테고리 없음

MLOps - BentoML 실습

Y's. 2024. 3. 17. 23:42

안녕하세요, 이번 포스팅에서는 실제 BentoML 을 통해서 모델 서빙이 어떻게 이루어지느지에 대해 알아보겠습니다.

이번 실습에 사용할 모델은 T2I 로 유명한 Stable Diffusion 모델을 사용해볼예정인데요!

아무래도 gpu 가 필요하다보니 클라우드 환경에서의 실습예제로 소개드리려합니다!

MLOps 토이 프로젝트 WorkFlow

이번 실습에서 사용된 모듈버전은 다음과 같습니다.

bentoml = 1.1.5

torch ==  2.0.1

transformers == 4.31.0

accelerate == 0.21.0

diffusers == 0.19.3

pydantic == 2.3.0

이번 실습의 단계를 구분하면 총 5 가지로 구분할 수 있는데요.

지금부터 하나 씩 살펴보겠습니다!

 

 

 

1. Saving a Model

첫번째로는 BentoML을 이용하여 BentoML 전용 model store에 모델을 저장하는 단계인데요.

이 단계는 모델 버전 관리 및 meta data를 같이 저장하는 단계입니다.

import bentoml

bentoml.diffusers.import_model(
    "sd-xl",  # Model tag in the BentoML Model Store
    "stabilityai/stable-diffusion-xl-base-1.0",  # Hugging Face model identifier
    signatures={
        "__call__": {"batchable": False},
        "text2img": {"batchable": False},
        "img2img": {"batchable": False},
        "inpaint": {"batchable": False},
    },
)
  • bentoml.diffusers.import_model을 사용하여 diffusers 기반인 stable-diffusion-xl 을 sd-xl 이라는 이름을 저장합니다.
  • signatures 파라미터인 batchable 을 False 로 하여 batch를1만 사용하도록 제한합니다.
    • batchable를 True로 할 경우, 모델의 추론요청을 할 때 batch를 input으로 받아서 추론이 가능합니다.

위 파일을 실행하면 bentoml/models/[model tag] 경로에 아래와 같은 모델 정보가 저장됩니다.

저장된 모델 정보

 

 

 

2. Definition Model Args

다음 단계에서는 Service 에서 사용하는 api 요청을 할 때, input의 type을 정의해보도록 하겠습니다.

type을 미리 정의해줌으로써 API 호출 시 제약조건을 보장시키기 위해 사용할 수 있습니다.

import typing as t

from pydantic import BaseModel

class SDArgs(BaseModel):
    prompt: t.Optional[str]
    negative_prompt: t.Optional[str] = "lowres, cropped, worst quality, blurry, bad art, bad proportios, simple background, duplicate"

    class Config:
        extra = "allow"
  • prompt: string 타입
  • negative_prompt: string 타입이며 미리 prompt를 정의
  • config: 추가 파라미터를 받을 수 있도록 allow

 

 

 

3. Creating a Service

Service는 BentoML의 core component 라고 이전 포스팅에서 설명드렸던는데 기억하시나요?

간단하게 말하면 model serving의 logic 을 담는 주체라고 생각하면 되는데요.

우리는 이 service를 통해 model inference를 하도록 하는 api endpoint를 쉽게 만들 수 있습니다.

여기서 우리는 앞선 단계에서 정의했던 Type을 사용하여 endpoint의 input 타입을 정의해주도록 하겠습니다.

import bentoml
from bentoml.io import JSON

from sdargs import SDArgs

import base64
from io import BytesIO

import torch

bento_model = bentoml.diffusers.get("sd-xl:latest")

sdxl_runner = bento_model.with_options(
    pipeline_class="diffusers.StableDiffusionXLPipeline",
    torch_dtype=torch.float16,
    variant="fp16",
    enable_torch_compile=True,
    enable_xformers=True,
).to_runner(name="sd-xl-runner")

svc = bentoml.Service(
    "stable-diffusion-xl",
    runners=[sdxl_runner],
)

@svc.api(input=JSON(pydantic_model=SDArgs), output=JSON())
def txt2img(input_data):
    kwargs = input_data.dict()
    results = sdxl_runner.run(**kwargs)

    images = result[0]
    img_str = []

    buffered = BytesIO()
    images.save(buffered, format="png")
    img_str.append(base64.b64encode(buffered.getvalue()).decode("utf-8"))

    return {"image": img_str}

이제 위 코드에 대해 하나씩 설명해드리도록 하겠습니다.

  1. 앞선 단계에서 저장한 sd-xl을 load하기 위해 bentoml.diffusers.get 함수를 사용합니다.
  2. to_runner 함수를 통해 모델을 실행(inference)할 수 있는 하나의 computation unit으로 만듭니다.
    이때, runner는 remote python worker에서 실행이되며 scaling 기능을 가지고 있습니다.
  3. bentoml.Service 함수를 통해 stable-diffusion-xl 이름의 service를 생성합니다.
    이 과정에서 우리는 Service를 handling 하는 svc 라는 변수로 선언합니다.
  4. @svc.api(input=JSON(pydantic_model=SDArgs), output=JSON()) 을 통해 svc service의 inference api endpoint를 만듭니다.
    한 가지 특이한 점은 BentoML의 경우 endpoint를 flask, fastAPI 처럼 데코레이터로 정의해주지 않고 함수명이 곧 endpoint를 의미합니다.
  5. txt2img 함수 내에서 inference를 위한 service logic을 구현합니다.

이제 우리는 ~/txt2img 라는 경로로 stable diffusion xl 모델을 추론하는 api를 정의했습니다.

 

 

 

4. Setting Config

api를 정의한 다음으로는 실제 bentoml을 사용해서 모델을 서빙할 때, 정의한 runner가 사용할 gpu를 지정하고 api server의 port 와 같은 config 값을 정의해주는 yaml 파일을 생성해보도록 하겠습니다.

runners:
    sd-xl-runner:
        resources:
            nvidia.com/gpu: [0]
api_server:
    http:
        port: 3002
    traffic:
        timeout: 5000
  • runners
    • sd-xl-runner: runners 중 sd-xl runner 의 이름을 가진 runner가 gpu 0 을 사용하도록 지정합니다.
  • api_server
    • http
      • port: 3002 port로 api를 serving 하도록 지정합니다.
    • traffic
      • timeout: response를 받기 전 대기할 최대 시간(5000초)을 지정합니다.

 

 

 

5. Serving Test

이제 우리는 다음의 단계를 통해서 모델을 서빙해볼 수 있습니다.

  1. diffusion 모델을 BentoML 모델 스토어로 가져오기
  2. BENTOML_CONFIG=./configuration.yaml bentoml serve service:svc 을 입력하여 serving 테스트 시작

service: service.py 를 의미

svc: service.py 내의 service 주체인 svc 변수를 의미

위와 같이 출력된다면 정상적으로 service가 실행 중임을 확인할 수 있습니다.

이제 실제 3002 port에 txt2img endpoint 로 요청을 보내면 응답을 받을 수 있습니다.

 

 

 

정말 간단하게 stable diffusion xl 모델을 BentoML 프레임워크를 사용해서 api 로 서빙하는 실습을 해보았는데요.

실제 BentoML 의 경우는 딥러닝 모델 배포를 위한 프레임워크 이다 보니 내부적으로 로드밸런싱을 비롯하여 큐작업까지 구현이 되어 있어서 실제 구현 방식이 복잡할 수 있는 대기큐와 로드밸런싱을 직접 구현해줄 필요가 없게 됩니다!

 

이는 요즘같이 AI 기술을 이용한 PoC 수행이 많아짐에 따라 효율적인 시간관리가 가능하게 해준다는 큰 장점이 있는데요.

요즘에는 BentoML 말고도 kserve, triton 등과 같은 여러 서빙 프레임워크들이 있다라고 해요.

 

이 밖에도 이러한 프레임워크의 장점으로는 MLFlow와 같은 Model tracking에 특화된 프레임워크와의 연동을 꼽아줄 수 있는데요. 실제 우리는 허깅페이스의 있는 stable diffusion 모델을 직접 다운로드해서 사용하는 방식을 선택했지만 Custom 한 모델이 있을 경우 이를 MLFlow 에 올려두어 모델을 저장한다면 BentoML로 직접 다운로드 해서 모델 실험 및 배포까지 만들 수 있는 하나의 파이프라인 구축을 가능하게 합니다. (우리가 최종적으로 만들 파이프라인이기도 합니다!)

 

나아가 BentoML 의 모니터링 부분은 Yatai 서버에 올려두어 실제 해당 모델의 배포상태에 따른 모니터링까지 확인해볼 수 있는데요. 이 부분은 차차 글로 정리하여 공유드리도록 하겠습니다!

 

 

 

 

 

 

 

References