Adding Custom Model Artifact¶
BentoML integrates with the most popular machine learning frameworks. For the ML framework yet to integrate with BentoML, BentoML provides model artifact customizing…
The guide will demonstrate how to create a custom model artifact class, and then use it in BentoService for prediction
1. Create custom Artifact¶
The following code creates a subclass from the BentoServiceArtifact. It implements how to save and load the model. In the pack method, the model class does validation to make sure the model is valid. It uses cloudpickle to save and load.
# my_model_artifact.py
import os
import json
from bentoml.utils import cloudpickle
from bentoml.exceptions import InvalidArgument
from bentoml.service.artifacts import BentoServiceArtifact
class MyModelArtifact(BentoServiceArtifact):
def __init__(self, name):
super(MyModelArtifact, self).__init__(name)
self._model = None
def pack(self, model, metadata=None):
if isinstance(model, dict) is not True:
raise InvalidArgument('MyModelArtifact only support dict')
if model.get('foo', None) is None:
raise KeyError('"foo" is not available in the model')
self._model = model
return self
def get(self):
return self._model
def save(self, dst):
with open(self._file_path(dst), 'wb') as file:
cloudpickle.dump(self._model, file)
def load(self, path):
with open(self._file_path(path), 'rb') as file:
model = cloudpickle.load(file)
return self.pack(model)
def _file_path(self, base_path):
return os.path.join(base_path, self.name + '.json')
2. Define and save BentoService with the custom Artifact¶
# my_bento_service.py
from my_model_artifact import MyModelArtifact
from bentoml import BentoService, env, api, artifacts
from bentoml.adapters import JsonInput
import bentoml
@env(infer_pip_packages=True)
@artifacts([MyModelArtifact('test_model')])
class MyService(bentoml.BentoService):
@api(input=JsonInput(), batch=False)
def predict(self, input_data):
result = input_data['bar'] + self.artifacts.test_model['foo']
return {'result': result}
from my_bento_service import MyService
svc = MyService()
model = {'foo': 2}
svc.pack('test_model', model)
svc.save()
3. Test with example data¶
$ bentoml serve MyService:latest
In another terminal to make a curl request
$ curl -i --header "Content-Type: application/json" \
--request POST --data '{"bar": 1}' \
http://localhost:5000/predict
# Output
HTTP/1.0 400 BAD REQUEST
X-Request-Id: cb63a61e-dc2a-4e12-a91c-8b15316a99df
Content-Type: text/html; charset=utf-8
Content-Length: 20
Server: Werkzeug/0.15.4 Python/3.7.3
Date: Tue, 16 Mar 2021 01:47:38 GMT
'{"result": 3}'%