BentoML 1.0에서 MLflow 모델 사용 방법

정보

• BentoML 1.0.7, MLflow 1.28.0 기준입니다.
• MLflow artifact를 로컬에 저장했다고 가정하고 설명하겠습니다.

예시 코드

BentoML 1.0에서 MLflow 모델을 사용할 때는 Runnable class를 만들어서 사용하면 됩니다.

import bentoml
from bentoml.io import JSON
import mlflow


class TestModelRunnable(bentoml.Runnable):
    # SUPPORTED_RESOURCES = ('cpu',)
    SUPPORTED_RESOURCES = ('nvidia.com/gpu', )
    SUPPORTS_CPU_MULTI_THREADING = True

    def __init__(self):
        self.model = mlflow.pytorch.load_model('test_model_dir')

    @bentoml.Runnable.method(batchable=False)
    def __call__(self, json_data):
        return self.model.predict(json_data)


runner = bentoml.Runner(TestModelRunnable, name='test_model')
svc = bentoml.Service('test_model_service', runners=[runner])

@svc.api(input=JSON(), output=JSON(), route='/test_model_predict')
async def predict(json_data):
    return await runner.async_run(json_data)

SUPPORTED_RESOURCES는 'cpu', 'nvidia.com/gpu' 둘 중 하나를 선택해서 사용 가능합니다. 사용하는 자원에 맞춰 선택하시면 됩니다.