Repository: bentoml/BentoML
Pull request #14 (2019-04-04 10:06:53)

Make BentoService.api decorator required for user model
What changes were proposed in this pull request?
------------------------------------------------
* remove "predict" magic, make api decorator required for user model

Also included a few other improvements
* fix module save and load in a few different scenarios 
* improve s3 save/load implementation
* skip s3 save/load test

Does this close any currently open issues?
------------------------------------------
n/a

How was this patch tested?
--------------------------
local tox passed


Patch:
diff --git a/bentoml/__init__.py b/bentoml/__init__.py
index d817dae1..aa83c3f7 100644
--- a/bentoml/__init__.py
+++ b/bentoml/__init__.py
@@ -18,10 +18,13 @@ from __future__ import absolute_import
 from __future__ import division
 from __future__ import print_function
 
+from bentoml import handlers
 from bentoml.server import metrics
 from bentoml.model import BentoModel
-from bentoml.service import BentoService, handler_decorator as handler
-from bentoml.version import __version__
+from bentoml.service import BentoService
 from bentoml.loader import load
+from bentoml.version import __version__
+
+api = BentoService.api
 
-__all__ = ['__version__', 'BentoModel', 'BentoService', 'load', 'handler', 'metrics']
+__all__ = ['__version__', 'api', 'BentoModel', 'BentoService', 'load', 'handlers', 'metrics']
diff --git a/bentoml/artifacts/pytorch_model_artifact.py b/bentoml/artifacts/pytorch_model_artifact.py
index f99d4d8c..d90a58b4 100644
--- a/bentoml/artifacts/pytorch_model_artifact.py
+++ b/bentoml/artifacts/pytorch_model_artifact.py
@@ -42,7 +42,11 @@ class PytorchModelArtifact(Artifact):
         return os.path.join(base_path, self.name + self._file_extension)
 
     def pack(self, model):  # pylint:disable=arguments-differ
-        import torch
+        try:
+            import torch
+        except ImportError:
+            raise ImportError("torch package is required to use PytorchModelArtifact")
+
         if not isinstance(model, torch.nn.Module):
             raise TypeError("PytorchModelArtifact can only pack type 'torch.nn.Module'")
 
diff --git a/bentoml/artifacts/tf_keras_model_artifact.py b/bentoml/artifacts/tf_keras_model_artifact.py
index 38758c42..de31d178 100644
--- a/bentoml/artifacts/tf_keras_model_artifact.py
+++ b/bentoml/artifacts/tf_keras_model_artifact.py
@@ -36,14 +36,17 @@ class TfKerasModelArtifact(Artifact):
     def model_file_path(self, base_path):
         return os.path.join(base_path, self.name + self._model_extension)
 
-    def pack(self, model):
+    def pack(self, model):  # pylint:disable=arguments-differ
         self.model = model
 
     def get(self):
         return self.model
 
-    def load(self, base_path):  # pylint:disable=arguments-differ
-        from tensorflow.keras.models import load_model
+    def load(self, base_path):
+        try:
+            from tensorflow.keras.models import load_model
+        except ImportError:
+            raise ImportError("tensorflow package is required to use TfKerasModelArtifact")
         self.model = load_model(self.model_file_path(base_path))
 
     def save(self, base_path):
diff --git a/bentoml/cli/__init__.py b/bentoml/cli/__init__.py
index 243bd1eb..0f3f4120 100644
--- a/bentoml/cli/__init__.py
+++ b/bentoml/cli/__init__.py
@@ -23,7 +23,7 @@ import os
 import tempfile
 from bentoml.loader import load
 from bentoml.server import BentoModelApiServer
-from bentoml.utils.s3 import is_s3_path, download_from_s3
+from bentoml.utils.s3 import is_s3_url, download_from_s3
 
 
 @click.group()
@@ -43,10 +43,10 @@ def serve(model_path, port):
     """
     port = port if port is not None else 5000
 
-    if is_s3_path(model_path):
+    if is_s3_url(model_path):
         temp_dir = tempfile.mkdtemp()
-        downloaded_file_path = download_from_s3(model_path, temp_dir)
-        model_path = downloaded_file_path
+        download_from_s3(model_path, temp_dir)
+        model_path = temp_dir
     else:
         if not os.path.isabs(model_path):
             model_path = os.path.abspath(model_path)
diff --git a/bentoml/loader.py b/bentoml/loader.py
index cc02e707..ff022a98 100644
--- a/bentoml/loader.py
+++ b/bentoml/loader.py
@@ -21,51 +21,49 @@ from __future__ import print_function
 import os
 import sys
 import tempfile
-import uuid
 
 from ruamel.yaml import YAML
 
 from bentoml.service import BentoService
 from bentoml.version import __version__ as BENTOML_VERSION
-from bentoml.utils.s3 import is_s3_path, download_from_s3
+from bentoml.utils.s3 import is_s3_url, download_from_s3
 
 
 class _LoadedBentoServiceWrapper(BentoService):
 
     def __init__(self, model_service_class, path, config):
         super(_LoadedBentoServiceWrapper, self).__init__()
-        self.path = path
-        self.config = config
-        self.model_service = model_service_class()
-        self.loaded = False
+        self._path = path
+        self._config = config
+        self._model_service = model_service_class()
         self._wrap_api_funcs()
+        self.loaded = False
 
     def load(self, path=None):
         if path is not None:
             # TODO: warn user path is ignored when using pip installed bentoML model
             pass
-        self.model_service.load(self.path)
+        self._model_service.load(self._path)
         self.loaded = True
 
-    @property
-    def apis(self):
-        return self.model_service._apis
+    def get_service_apis(self):
+        return self._model_service.get_service_apis()
 
     @property
     def name(self):
-        return self.config['model_name']
+        return self._config['model_name']
 
     @property
     def version(self):
-        return self.config['model_version']
+        return self._config['model_version']
 
     def _wrap_api_funcs(self):
         """
         Add target ModelService's API methods to the returned wrapper class
         """
-        for api in self.apis:
+        for api in self.get_service_apis():
             setattr(self, api.name,
-                    self.model_service.__getattribute__(api.name).__get__(self.model_service))
+                    self._model_service.__getattribute__(api.name).__get__(self._model_service))
 
 
 def load_bentoml_config(path):
@@ -96,36 +94,36 @@ def load(path, lazy_load=False):
     :param path: A BentoArchive path generated from BentoService.save call
     :return: BentoService
     """
-    if is_s3_path(path):
-        temp_dir = tempfile.mkdtemp()
-        downloaded_file_path = download_from_s3(path, temp_dir)
-
-        file_path = downloaded_file_path
-    else:
-        file_path = path
+    if is_s3_url(path):
+        tempdir = tempfile.mkdtemp()
+        download_from_s3(path, tempdir)
+        path = tempdir
 
-    config = load_bentoml_config(file_path)
+    config = load_bentoml_config(path)
 
     # Load target module containing BentoService class from given path
-    module_file_path = os.path.join(file_path, config['model_name'], config['module_file'])
+    module_file_path = os.path.join(path, config['model_name'], config['module_file'])
 
-    if sys.version_info >= (3, 5):
+    module_name = config['module_name']
+    if module_name in sys.modules:
+        # module already loaded, TODO: add warning
+        module = sys.modules[module_name]
+    elif sys.version_info >= (3, 5):
         import importlib.util
-        spec = importlib.util.spec_from_file_location(config['module_name'], module_file_path)
+        spec = importlib.util.spec_from_file_location(module_name, module_file_path)
         module = importlib.util.module_from_spec(spec)
         spec.loader.exec_module(module)
     elif sys.version_info >= (3, 3):
         from importlib.machinery import SourceFileLoader
         # pylint:disable=deprecated-method
-        module = SourceFileLoader(config['module_name'],
-                                  module_file_path).load_module(config['module_name'])
+        module = SourceFileLoader(module_name, module_file_path).load_module(module_name)
         # pylint:enable=deprecated-method
     else:
         import imp
-        module = imp.load_source(config['module_name'], module_file_path)
+        module = imp.load_source(module_name, module_file_path)
 
     model_service_class = module.__getattribute__(config['model_name'])
-    loaded_model = _LoadedBentoServiceWrapper(model_service_class, file_path, config)
+    loaded_model = _LoadedBentoServiceWrapper(model_service_class, path, config)
 
     if not lazy_load:
         loaded_model.load()
diff --git a/bentoml/model.py b/bentoml/model.py
index 6da9da9a..78cad746 100644
--- a/bentoml/model.py
+++ b/bentoml/model.py
@@ -27,9 +27,10 @@ from datetime import datetime
 from abc import abstractmethod
 from bentoml.utils import Path
 from bentoml.utils.module_helper import copy_module_and_local_dependencies
-from bentoml.utils.s3 import is_s3_path, upload_to_s3, download_from_s3
+from bentoml.utils.exceptions import BentoMLException
+from bentoml.utils.s3 import is_s3_url, upload_to_s3, download_from_s3
 from bentoml.service_env import BentoServiceEnv
-from bentoml.service import SingleModelBentoService
+from bentoml.service import BentoService
 from bentoml.artifacts import ArtifactCollection
 from bentoml.loader import load_bentoml_config
 from bentoml.version import __version__ as BENTOML_VERSION
@@ -169,7 +170,7 @@ def _generate_new_version_str():
     return date_string + '_' + random_hash
 
 
-class BentoModel(SingleModelBentoService):
+class BentoModel(BentoService):
     """
     BentoModel is the base abstraction for defining how a ML model can
     be save to or load from a file directory in BentoArchive format
@@ -221,7 +222,7 @@ class BentoModel(SingleModelBentoService):
             version = _generate_new_version_str()
         self._version = version
 
-        if is_s3_path(base_path):
+        if is_s3_url(base_path):
             storage_type = 's3'
             temp_dir = tempfile.mkdtemp()
             remote_path = base_path
@@ -316,19 +317,20 @@ class BentoModel(SingleModelBentoService):
             path = self.__class__._bento_module_path
             artifacts_path = path
         else:
+            if path is None:
+                raise BentoMLException("Loading path is required for {}#load".format(self.name))
+
             # When calling load on generated archive directory, look for /artifacts
             # directory under module sub-directory
+            if is_s3_url(path):
+                temporary_path = tempfile.mkdtemp()
+                download_from_s3(path, temporary_path)
+                # Use loacl temp path for the following loading operations
+                path = temporary_path
 
-            # TODO: assert path is not None
+            artifacts_path = os.path.join(path, self.name)
 
-            if is_s3_path(path):
-                temporary_path = tempfile.mkdtemp()
-                download_file_path = download_from_s3(path, temporary_path)
-                artifacts_path = os.path.join(download_file_path, self.name)
-                bentoml_config = load_bentoml_config(download_file_path)
-            else:
-                artifacts_path = os.path.join(path, self.name)
-                bentoml_config = load_bentoml_config(path)
+        bentoml_config = load_bentoml_config(path)
 
         self.artifacts.load(artifacts_path)
         self._model_name = bentoml_config['model_name']
@@ -339,7 +341,3 @@ class BentoModel(SingleModelBentoService):
     @abstractmethod
     def config(self, artifacts, env):
         pass
-
-    @abstractmethod
-    def predict(self, data):
-        pass
diff --git a/bentoml/server/model_api_server.py b/bentoml/server/model_api_server.py
index 4caa4123..26eb6c97 100644
--- a/bentoml/server/model_api_server.py
+++ b/bentoml/server/model_api_server.py
@@ -143,7 +143,7 @@ class BentoModelApiServer():
         model_name/model_version/api_name
         """
         model_service = self.model_service
-        apis = self.model_service.apis
+        apis = self.model_service.get_service_apis()
         if len(apis) == 1:
             self.setup_api_func_route(model_service.name, model_service.version, apis[0].name,
                                       apis[0])
diff --git a/bentoml/service.py b/bentoml/service.py
index 2b290e5b..b95b4335 100644
--- a/bentoml/service.py
+++ b/bentoml/service.py
@@ -45,9 +45,6 @@ class BentoService(object):
     for BentoAPIServer and BentoCLI
     """
 
-    def __init__(self):
-        self._config_service_apis()
-
     @abstractmethod
     def load(self, path):
         """
@@ -55,14 +52,14 @@ class BentoService(object):
         """
 
     def _config_service_apis(self):
-        self._apis = []
+        self._service_apis = []   # pylint:disable=attribute-defined-outside-init
         for _, function in inspect.getmembers(
                 self.__class__, predicate=lambda x: inspect.isfunction(x) or inspect.ismethod(x)):
             if hasattr(function, '_is_api'):
                 api_name = _get_func_attr(function, '_api_name')
                 handler = _get_func_attr(function, '_handler')
                 func = function.__get__(self)
-                self._apis.append(BentoServiceAPI(api_name, handler, func))
+                self._service_apis.append(BentoServiceAPI(api_name, handler, func))
 
     @property
     @abstractmethod
@@ -74,9 +71,10 @@ class BentoService(object):
     def version(self):
         pass
 
-    @property
-    def apis(self):
-        return self._apis
+    def get_service_apis(self):
+        if not hasattr(self, '_service_apis'):
+            self._config_service_apis()
+        return self._service_apis
 
     @staticmethod
     def api(handler=DataframeHandler, api_name=None):
@@ -140,42 +138,3 @@ class BentoServiceAPI(object):
     @property
     def func(self):
         return self._func
-
-
-class SingleModelBentoService(BentoService):
-    """
-    A BentoService that host a single model's predict function as its only API
-    """
-
-    def __init__(self):
-        # Adding BentoServiceAPI metadata to predict function
-        try:
-            handler = _get_func_attr(self.__class__.predict, "_handler")
-        except AttributeError:
-            handler = DataframeHandler
-        BentoService.api(handler, 'predict')(self.__class__.predict)
-
-        super(SingleModelBentoService, self).__init__()
-
-    @abstractmethod
-    def load(self, path):
-        pass
-
-    @abstractmethod
-    def predict(self, data):
-        pass
-
-
-def handler_decorator(handler=DataframeHandler):
-    """
-    Decorator for setting handler Class for predict function in a BentoModel
-    """
-
-    def _handler_decorator(predict_func):
-        if predict_func.__name__ != 'predict':
-            raise ValueError("Bentoml's handler decoration can only "
-                             "be applied to BentoModel#predict function")
-        _set_func_attr(predict_func, '_handler', handler)
-        return predict_func
-
-    return _handler_decorator
diff --git a/bentoml/utils/module_helper.py b/bentoml/utils/module_helper.py
index e3ee9c09..dcb740c2 100644
--- a/bentoml/utils/module_helper.py
+++ b/bentoml/utils/module_helper.py
@@ -36,7 +36,11 @@ def _get_module_src_file(module):
 
 def copy_module_and_local_dependencies(target_module, destination, toplevel_package_path=None,
                                        copy_entire_package=False):
-    """bundle given module, and all its dependencies, copy files to destination"""
+    """
+    bundle given module, and all its dependencies within top level package,
+    and copy all source files to destination path, essentially creating
+    a source distribution of target_module
+    """
     if isinstance(target_module, string_types):
         target_module = importlib.import_module(target_module)
     else:
@@ -54,6 +58,9 @@ def copy_module_and_local_dependencies(target_module, destination, toplevel_pack
 
     target_module_file = _get_module_src_file(target_module)
 
+    if target_module_name == '__main__':
+        target_module_name = target_module_file[:-3].replace(os.sep, '.')
+
     if copy_entire_package:
         if toplevel_package_path is None:
             raise BentoMLException("Must set toplevel_package_path when using"
@@ -100,25 +107,30 @@ def copy_module_and_local_dependencies(target_module, destination, toplevel_pack
                 if name not in local_modules and _get_module_src_file(module).startswith(path):
                     local_modules[name] = module
 
-    # Lastly, add target module itself
-    local_modules[target_module_name] = target_module
-
     # Remove "__main__" module, if target module is loaded as __main__, it should
     # be in module_files as (module_name, module_file) in current context
     if '__main__' in local_modules:
         del local_modules['__main__']
 
+    # Lastly, add target module itself
+    local_modules[target_module_name] = target_module
+
     for module_name, module in iteritems(local_modules):
         module_file = _get_module_src_file(module)
 
         with open(module_file, "rb") as f:
             src_code = f.read()
 
-        _, file_name = os.path.split(module_file)
-
-        if file_name == '__init__.py':
+        if not os.path.isabs(module_file):
+            # For modules within current top level package, module_file here should be a
+            # relative path to the src file
+            target_file = os.path.join(destination, module_file)
+        elif os.path.split(module_file)[1] == '__init__.py':
+            # for module a.b.c in 'some_path/a/b/c/__init__.py', copy file to
+            # 'destination/a/b/c/__init__.py'
             target_file = os.path.join(destination, module_name.replace('.', '/'), '__init__.py')
         else:
+            # for module a.b.c in 'some_path/a/b/c.py', copy file to 'destination/a/b/c.py'
             target_file = os.path.join(destination, module_name.replace('.', '/') + '.py')
 
         target_path = os.path.dirname(target_file)
diff --git a/bentoml/utils/s3.py b/bentoml/utils/s3.py
index e89ee09c..6fd096d1 100644
--- a/bentoml/utils/s3.py
+++ b/bentoml/utils/s3.py
@@ -22,93 +22,51 @@ import os
 import boto3
 from six.moves.urllib.parse import urlparse
 
+from bentoml.utils import Path
 
-def is_s3_path(path):
+def is_s3_url(url):
     """
-    This helper function return True if path is in s3:// format
-    """
-    parsed_url = urlparse(path)
-    if parsed_url.scheme == 's3':
-        return True
-    return False
-
-
-def find_bucket_key(s3_path):
-    """
-    This is a helper function that given an s3 path such that the path is of
-    the form: bucket/key
-    It will return the bucket and the key represented by the s3 path
-    """
-    s3_components = s3_path.split('/')
-    bucket = s3_components[0]
-    s3_key = ""
-    if len(s3_components) > 1:
-        s3_key = '/'.join(s3_components[1:])
-    return bucket, s3_key
-
-
-def split_s3_bucket_key(s3_path):
-    """Split s3 path into bucket and key prefix.
-    This will also handle the s3:// prefix.
-    :return: Tuple of ('bucketname', 'keyname')
-    """
-    if s3_path.startswith('s3://'):
-        s3_path = s3_path[5:]
-    return find_bucket_key(s3_path)
-
-
-def create_bucket_if_not_exit(client, bucket):
-    """
-    Create s3 bucket with the giving bucket name, if it didn't exist
+    Check if url is an s3, s3n, or s3a url
     """
     try:
-        client.head_bucket(Bucket=bucket)
+        return urlparse(url).scheme in ['s3', 's3n', 's3a']
     except Exception:
-        region_name = boto3.Session().region_name
-        client.create_bucket(Bucket=bucket,
-                             CreateBucketConfiguration={'LocationConstraint': region_name})
+        return False
 
 
-def upload_to_s3(s3_path, file_path):
+def upload_to_s3(s3_url, file_path):
     """
     Update files in the file_path to the s3 location
     """
-    bucket, s3_key = split_s3_bucket_key(s3_path)
 
-    try:
-        s3_client = boto3.client('s3')
-        create_bucket_if_not_exit(s3_client, bucket)
+    parse_result = urlparse(s3_url)
+    bucket = parse_result.netloc
+    base_path = parse_result.path
 
-        for root, dirs, files in os.walk(file_path):
-            for file_name in files:
-                submit_file_path = os.path.join(root, file_name)
-                full_path = submit_file_path[len(file_path) + 1:]
-                s3_path = os.path.join(s3_key, full_path)
-                s3_client.upload_file(Filename=submit_file_path, Bucket=bucket, Key=s3_path)
+    s3_client = boto3.client('s3')
 
-        #TODO: Clean up files in the temp dir
-    except Exception as e:
-        raise e
+    for root, _, files in os.walk(file_path):
+        for file_name in files:
+            abs_file_path = os.path.join(root, file_name)
+            relative_file_path = abs_file_path[len(file_path) + 1:]
+            s3_path = os.path.join(base_path, relative_file_path)
+            s3_client.upload_file(Filename=abs_file_path, Bucket=bucket, Key=s3_path)
 
 
-def download_from_s3(s3_path, file_path):
+def download_from_s3(s3_url, file_path):
     """
     Download files from given s3_path and store in the given file path
     """
-    bucket, s3_key = split_s3_bucket_key(s3_path)
-
-    try:
-        s3_client = boto3.client('s3')
-        list_object_result = s3_client.list_objects(Bucket=bucket, Prefix=s3_key)
-        result_content = list_object_result['Contents']
-
-        for content in result_content:
-            full_path = content['Key'][len(s3_key) + 1:]
-            local_file_path = os.path.join(file_path, full_path)
-            if not os.path.exists(os.path.dirname(local_file_path)):
-                os.makedirs(os.path.dirname(local_file_path))
-            s3_client.download_file(Bucket=bucket, Key=content['Key'], Filename=local_file_path)
-
-        return file_path
-    except Exception as e:
-        raise e
+    parse_result = urlparse(s3_url)
+    bucket = parse_result.netloc
+    base_path = parse_result.path
+
+    s3_client = boto3.client('s3')
+    list_object_result = s3_client.list_objects(Bucket=bucket, Prefix=base_path)
+    result_content = list_object_result['Contents']
+
+    for content in result_content:
+        relative_file_path = content['Key'][len(base_path) + 1:]
+        local_file_path = os.path.join(file_path, relative_file_path)
+        Path(os.path.dirname(local_file_path)).mkdir(parents=True, exist_ok=True)
+        s3_client.download_file(Bucket=bucket, Key=content['Key'], Filename=local_file_path)
\ No newline at end of file
diff --git a/examples/local-dependencies/mymodule/scripts/main.py b/examples/local-dependencies/mymodule/scripts/main.py
index ff9ef959..87936d35 100644
--- a/examples/local-dependencies/mymodule/scripts/main.py
+++ b/examples/local-dependencies/mymodule/scripts/main.py
@@ -5,10 +5,11 @@ from sklearn import datasets
 
 # Use local bentoml code
 sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), '../../../..')))
-from bentoml import BentoModel, handler, load
+from bentoml import BentoModel, api, load
 from bentoml.artifacts import PickleArtifact
 from bentoml.handlers import JsonHandler
 
+sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), '../..')))
 from mymodule import method_in_mymodule
 from mymodule.submodule import method_in_submodule
 from mymodule.submodule1 import method_in_submodule1
@@ -23,7 +24,7 @@ class IrisClassifier(BentoModel):
         artifacts.add(PickleArtifact('clf'))
         env.add_conda_dependencies(["scikit-learn"])
 
-    @handler(JsonHandler)
+    @api(JsonHandler)
     def predict(self, parsed_json):
         data = method_in_mymodule(parsed_json)
         data = method_in_submodule(data)
diff --git a/examples/quick-start/main.py b/examples/quick-start/main.py
index c0b454c8..9a1efdc2 100644
--- a/examples/quick-start/main.py
+++ b/examples/quick-start/main.py
@@ -3,7 +3,9 @@ import sys
 from sklearn import svm
 from sklearn import datasets
 
-from bentoml import BentoModel, handler, load
+# Use local bentoml code
+sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), '../..')))
+from bentoml import BentoModel, load, api
 from bentoml.artifacts import PickleArtifact
 from bentoml.handlers import JsonHandler
 
@@ -16,7 +18,7 @@ class IrisClassifier(BentoModel):
         artifacts.add(PickleArtifact('clf'))
         env.add_conda_dependencies(["scikit-learn"])
 
-    @handler(JsonHandler)
+    @api(JsonHandler)
     def predict(self, parsed_json):
         return self.artifacts.clf.predict(parsed_json)
 
@@ -27,8 +29,10 @@ if __name__ == "__main__":
     clf.fit(X, y)
 
     model = IrisClassifier(clf=clf)
+    print(X[0:1])
     print("model.predict = {}".format(model.predict(X[0:1])))
 
+    print("Saving model as bento archive...")
     saved_path = model.save("./model")
     print("BentoML model archive saved to path: {}".format(saved_path))
 


Test patch:
diff --git a/tests/server/test_model_api_server.py b/tests/server/test_model_api_server.py
index 944336df..92aebb7a 100644
--- a/tests/server/test_model_api_server.py
+++ b/tests/server/test_model_api_server.py
@@ -23,6 +23,7 @@ class RestServiceTestModel(bentoml.BentoModel):
 
         env.add_conda_dependencies(["scikit-learn"])
 
+    @bentoml.api(bentoml.handlers.DataframeHandler)
     def predict(self, df):
         """
         predict expects dataframe as input
diff --git a/tests/test_model.py b/tests/test_model.py
index 65f6ef01..9877d863 100644
--- a/tests/test_model.py
+++ b/tests/test_model.py
@@ -27,6 +27,7 @@ class MyFakeBentoModel(bentoml.BentoModel):
     def config(self, artifacts, env):
         artifacts.add(bentoml.artifacts.PickleArtifact('fake_model'))
 
+    @bentoml.api(bentoml.handlers.DataframeHandler)
     def predict(self, df):
         """
         predict expects dataframe as input
@@ -53,8 +54,8 @@ def test_save_and_load_model():
     model_service.load()
     assert model_service.loaded
 
-    assert len(model_service.apis) == 1
-    api = model_service.apis[0]
+    assert len(model_service.get_service_apis()) == 1
+    api = model_service.get_service_apis()[0]
     assert api.name == 'predict'
     assert api.handler == bentoml.handlers.DataframeHandler
     assert api.func(1) == 2
@@ -62,13 +63,11 @@ def test_save_and_load_model():
     # Check api methods are available
     assert model_service.predict(1) == 2
 
-
+@pytest.mark.skip(reason="Setup s3 creds in travis or use a mock")
 def test_save_and_load_model_from_s3():
     fake_model = MyFakeModel()
     sm = MyFakeBentoModel(fake_model=fake_model)
 
-    import boto3
-
     s3_location = 's3://bentoml/test'
     s3_saved_path = sm.save(base_path=s3_location)
     assert s3_saved_path == os.path.join(s3_location, sm.name, sm.version)
@@ -77,4 +76,4 @@ def test_save_and_load_model_from_s3():
     assert not download_model_service.loaded
     download_model_service.load()
     assert download_model_service.loaded
-    assert download_model_service.apis[0].func(1) == 2
+    assert download_model_service.get_service_apis()[0].func(1) == 2