Skip to content
Permalink
Browse files

Adding SklearnModelArtifact for save/load sklearn models (#308)

* rename - ArtifactInstance to ArtifactWrapper

* h2o_model_artifact.py

* xgboost_model_artifact.py

* internal renaming - ArtifactSpec to BentoServiceArtifact

* Create sklearn_model_artifact.py

* update sklearn model artifact

* update quick-start-guide to use sklearnmodelartifact

* Update README.md
  • Loading branch information...
parano committed Sep 24, 2019
1 parent ebde03b commit ab75076033ef05b44fcb9b2ce069479c5717416c
@@ -45,14 +45,14 @@ Defining a machine learning service with BentoML:

```python
import bentoml
from bentoml.artifact import PickleArtifact
from bentoml.artifact import SklearnModelArtifact
from bentoml.handlers import DataframeHandler
# You can also import your own Python module here and BentoML will automatically
# figure out the dependency chain and package all those Python modules
import my_preproceesing_lib
@bentoml.artifacts([PickleArtifact('model')])
@bentoml.artifacts([SklearnModelArtifact('model')])
@bentoml.env(pip_dependencies=["scikit-learn"])
class IrisClassifier(bentoml.BentoService):
@@ -79,7 +79,10 @@ X, y = iris.data, iris.target
clf.fit(X, y)
# Packaging trained model for serving in production:
saved_path = IrisClassifier.pack(model=clf).save('/tmp/bento')
iris_classifier_service = IrisClassifier.pack(model=clf)
# Save prediction service to file archive
saved_path = = iris_classifier_service.save()
```

A Bento is a versioned archive, containing the BentoService you defined, along
@@ -16,18 +16,23 @@
from __future__ import division
from __future__ import print_function

from bentoml.artifact.artifact import ArtifactSpec, ArtifactInstance, ArtifactCollection
from bentoml.artifact.artifact import (
BentoServiceArtifact,
BentoServiceArtifactWrapper,
ArtifactCollection,
)
from bentoml.artifact.text_file_artifact import TextFileArtifact
from bentoml.artifact.pickle_artifact import PickleArtifact
from bentoml.artifact.pytorch_model_artifact import PytorchModelArtifact
from bentoml.artifact.text_file_artifact import TextFileArtifact
from bentoml.artifact.keras_model_artifact import KerasModelArtifact
from bentoml.artifact.xgboost_artifact import XgboostModelArtifact
from bentoml.artifact.h2o_artifact import H2oModelArtifact
from bentoml.artifact.xgboost_model_artifact import XgboostModelArtifact
from bentoml.artifact.h2o_model_artifact import H2oModelArtifact
from bentoml.artifact.fastai_model_artifact import FastaiModelArtifact
from bentoml.artifact.sklearn_model_artifact import SklearnModelArtifact

__all__ = [
"ArtifactSpec",
"ArtifactInstance",
"BentoServiceArtifact",
"BentoServiceArtifactWrapper",
"ArtifactCollection",
"PickleArtifact",
"PytorchModelArtifact",
@@ -36,4 +41,5 @@
"XgboostModelArtifact",
"H2oModelArtifact",
"FastaiModelArtifact",
"SklearnModelArtifact",
]
@@ -21,12 +21,12 @@
ARTIFACTS_SUBPATH = "artifacts"


class ArtifactSpec(object):
class BentoServiceArtifact(object):
"""
Artifact is a spec describing how to pack and load different types
of model dependencies to/from file system.
A call to pack and load should return an ArtifactInstance that can
A call to pack and load should return an BentoServiceArtifactWrapper that can
be saved to file system or retrieved for BentoService workload
"""

@@ -56,10 +56,10 @@ def load(self, path):
"""


class ArtifactInstance(object):
class BentoServiceArtifactWrapper(object):
"""
ArtifactInstance is an object representing a materialized Artifact, either
loaded from file system or packed with data in a python session
BentoServiceArtifactWrapper is an object representing a materialized Artifact,
either loaded from file system or packed with data in a python session
"""

def __init__(self, spec):
@@ -68,7 +68,8 @@ def __init__(self, spec):
@property
def spec(self):
"""
:return: reference to the ArtifactSpec that produced this ArtifactInstance
:return: reference to the BentoServiceArtifact that generated this
BentoServiceArtifactWrapper
"""
return self._spec

@@ -79,8 +80,7 @@ def save(self, dst):

def get(self):
"""
Get returns an python object which provides all the functionality this
artifact provides
Get returns a reference to the artifact being packed
"""


@@ -103,11 +103,11 @@ def __getattr__(self, item):
return self[item].get()

def add(self, artifact):
if not isinstance(artifact, ArtifactInstance):
if not isinstance(artifact, BentoServiceArtifactWrapper):
raise TypeError(
"ArtifactCollection only accepts type bentoml.ArtifactInstance,"
"Must call Artifact#pack or Artifact#load before adding to"
"an ArtifactCollection"
"ArtifactCollection only accepts type BentoServiceArtifactWrapper,"
"Must call BentoServiceArtifact#pack or BentoServiceArtifact#load "
"before adding to an ArtifactCollection"
)

super(ArtifactCollection, self).__setitem__(artifact.spec.name, artifact)
@@ -124,7 +124,7 @@ def save(self, dst):
@classmethod
def load(cls, path, artifacts_spec):
"""bulk operation for loading all artifacts from path based on a list of
ArtifactSpec
BentoServiceArtifact
"""
load_path = os.path.join(path, ARTIFACTS_SUBPATH)
artifacts = cls()
@@ -20,10 +20,10 @@
import sys
import shutil

from bentoml.artifact import ArtifactSpec, ArtifactInstance
from bentoml.artifact import BentoServiceArtifact, BentoServiceArtifactWrapper


class FastaiModelArtifact(ArtifactSpec):
class FastaiModelArtifact(BentoServiceArtifact):
"""Saving and Loading FastAI Model
Args:
@@ -42,7 +42,7 @@ def _model_file_path(self, base_path):
return os.path.join(base_path, self._file_name)

def pack(self, model): # pylint:disable=arguments-differ
return _FastaiModelArtifactInstance(self, model)
return _FastaiModelArtifactWrapper(self, model)

def load(self, path):
try:
@@ -55,9 +55,9 @@ def load(self, path):
return self.pack(model)


class _FastaiModelArtifactInstance(ArtifactInstance):
class _FastaiModelArtifactWrapper(BentoServiceArtifactWrapper):
def __init__(self, spec, model):
super(_FastaiModelArtifactInstance, self).__init__(spec)
super(_FastaiModelArtifactWrapper, self).__init__(spec)
if sys.version_info.major < 3 or sys.version_info.minor < 6:
raise SystemError("fast ai requires python 3.6 version or higher")

@@ -19,10 +19,10 @@
import os
import shutil

from bentoml.artifact import ArtifactSpec, ArtifactInstance
from bentoml.artifact import BentoServiceArtifact, BentoServiceArtifactWrapper


class H2oModelArtifact(ArtifactSpec):
class H2oModelArtifact(BentoServiceArtifact):
"""Abstraction for saving/loading objects with h2o.save_model and h2o.load_model
Args:
@@ -36,7 +36,7 @@ def _model_file_path(self, base_path):
return os.path.join(base_path, self.name)

def pack(self, model): # pylint:disable=arguments-differ
return _H2oModelArtifactInstance(self, model)
return _H2oModelArtifactWrapper(self, model)

def load(self, path):
try:
@@ -49,9 +49,9 @@ def load(self, path):
return self.pack(model)


class _H2oModelArtifactInstance(ArtifactInstance):
class _H2oModelArtifactWrapper(BentoServiceArtifactWrapper):
def __init__(self, spec, model):
super(_H2oModelArtifactInstance, self).__init__(spec)
super(_H2oModelArtifactWrapper, self).__init__(spec)
self._model = model

def save(self, dst):
@@ -19,7 +19,7 @@
import os

from bentoml.utils import cloudpickle
from bentoml.artifact import ArtifactSpec, ArtifactInstance
from bentoml.artifact import BentoServiceArtifact, BentoServiceArtifactWrapper

try:
import tensorflow as tf
@@ -29,7 +29,7 @@
keras = None


class KerasModelArtifact(ArtifactSpec):
class KerasModelArtifact(BentoServiceArtifact):
"""
Abstraction for saving/loading Keras model
"""
@@ -94,7 +94,7 @@ def pack(self, data): # pylint:disable=arguments-differ

self.bind_keras_backend_session()
model._make_predict_function()
return _TfKerasModelArtifactInstance(self, model, custom_objects)
return _TfKerasModelArtifactWrapper(self, model, custom_objects)

def load(self, path):
if tf is None:
@@ -119,9 +119,9 @@ def load(self, path):
return self.pack(model)


class _TfKerasModelArtifactInstance(ArtifactInstance):
class _TfKerasModelArtifactWrapper(BentoServiceArtifactWrapper):
def __init__(self, spec, model, custom_objects):
super(_TfKerasModelArtifactInstance, self).__init__(spec)
super(_TfKerasModelArtifactWrapper, self).__init__(spec)

if tf is None:
raise ImportError(
@@ -21,10 +21,10 @@
import dill
from six import string_types

from bentoml.artifact import ArtifactSpec, ArtifactInstance
from bentoml.artifact import BentoServiceArtifact, BentoServiceArtifactWrapper


class PickleArtifact(ArtifactSpec):
class PickleArtifact(BentoServiceArtifact):
"""Abstraction for saving/loading python objects with pickle serialization
Args:
@@ -48,17 +48,17 @@ def _pkl_file_path(self, base_path):
return os.path.join(base_path, self.name + self._pickle_extension)

def pack(self, obj): # pylint:disable=arguments-differ
return _PickleArtifactInstance(self, obj)
return _PickleArtifactWrapper(self, obj)

def load(self, path):
with open(self._pkl_file_path(path), "rb") as pkl_file:
obj = self._pickle.load(pkl_file)
return self.pack(obj)


class _PickleArtifactInstance(ArtifactInstance):
class _PickleArtifactWrapper(BentoServiceArtifactWrapper):
def __init__(self, spec, obj):
super(_PickleArtifactInstance, self).__init__(spec)
super(_PickleArtifactWrapper, self).__init__(spec)

self._obj = obj

@@ -18,11 +18,11 @@

import os

from bentoml.artifact import ArtifactSpec, ArtifactInstance
from bentoml.artifact import BentoServiceArtifact, BentoServiceArtifactWrapper
from bentoml.utils import cloudpickle


class PytorchModelArtifact(ArtifactSpec):
class PytorchModelArtifact(BentoServiceArtifact):
"""
Abstraction for saving/loading objects with torch.save and torch.load
"""
@@ -35,7 +35,7 @@ def _file_path(self, base_path):
return os.path.join(base_path, self.name + self._file_extension)

def pack(self, model): # pylint:disable=arguments-differ
return _PytorchModelArtifactInstance(self, model)
return _PytorchModelArtifactWrapper(self, model)

def load(self, path):
try:
@@ -54,9 +54,9 @@ def load(self, path):
return self.pack(model)


class _PytorchModelArtifactInstance(ArtifactInstance):
class _PytorchModelArtifactWrapper(BentoServiceArtifactWrapper):
def __init__(self, spec, model):
super(_PytorchModelArtifactInstance, self).__init__(spec)
super(_PytorchModelArtifactWrapper, self).__init__(spec)

try:
import torch
@@ -0,0 +1,81 @@
# Copyright 2019 Atalaya Tech, Inc.

# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at

# http://www.apache.org/licenses/LICENSE-2.0

# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

from __future__ import absolute_import
from __future__ import division
from __future__ import print_function

import os

from bentoml.artifact import BentoServiceArtifact, BentoServiceArtifactWrapper

try:
import joblib
except ImportError:
joblib = None

if joblib is None:
try:
from sklearn.externals import joblib
except ImportError:
pass


class SklearnModelArtifact(BentoServiceArtifact):
"""
Abstraction for saving/loading scikit learn models using sklearn.externals.joblib
Args:
name (str): Name for the artifact
pickle_extension (str): The extension format for pickled file.
"""

def __init__(self, name, pickle_extension=".pkl"):
super(SklearnModelArtifact, self).__init__(name)

self._pickle_extension = pickle_extension

def _model_file_path(self, base_path):
return os.path.join(base_path, self.name + self._pickle_extension)

def pack(self, sklearn_model): # pylint:disable=arguments-differ
return _SklearnModelArtifactWrapper(self, sklearn_model)

def load(self, path):
if joblib is None:
raise ImportError(
"scikit-learn package is required to use " "SklearnModelArtifact"
)

model_file_path = self._model_file_path(path)
sklearn_model = joblib.load(model_file_path, mmap_mode='r')
return self.pack(sklearn_model)


class _SklearnModelArtifactWrapper(BentoServiceArtifactWrapper):
def __init__(self, spec, model):
super(_SklearnModelArtifactWrapper, self).__init__(spec)

self._model = model

def get(self):
return self._model

def save(self, dst):
if joblib is None:
raise ImportError(
"scikit-learn package is required to use " "SklearnModelArtifact"
)

joblib.dump(self._model, self.spec._model_file_path(dst))

0 comments on commit ab75076

Please sign in to comment.
You can’t perform that action at this time.