Repository: bentoml/BentoML
Pull request #18 (2019-04-05 21:17:47)

Add additional options for handlers
What changes were proposed in this pull request?
------------------------------------------------
Add configuration options to handlers. Solving issue such as https://github.com/bentoml/BentoML/issues/8. 


Does this close any currently open issues?
------------------------------------------
https://github.com/bentoml/BentoML/issues/8

How was this patch tested?
--------------------------



Patch:
diff --git a/bentoml/cli/commands.py b/bentoml/cli/commands.py
index 4d002263..b0b1c647 100644
--- a/bentoml/cli/commands.py
+++ b/bentoml/cli/commands.py
@@ -67,7 +67,7 @@ def run(ctx, model_path, api_name):
     else:
         matched_api = service_apis[matched_api_index]
 
-    matched_api.handler.handle_cli(ctx.args, matched_api.func)
+    matched_api.handler.handle_cli(ctx.args, matched_api.func, matched_api.options)
 
 
 # bentoml serve --model=/MODEL_PATH --port=PORT_NUM
diff --git a/bentoml/handlers/base_handlers.py b/bentoml/handlers/base_handlers.py
index 735dd620..879cd838 100644
--- a/bentoml/handlers/base_handlers.py
+++ b/bentoml/handlers/base_handlers.py
@@ -27,7 +27,7 @@ class RequestHandler():
 
     @staticmethod
     @abstractmethod
-    def handle_request(request, func):
+    def handle_request(request, func, options):
         pass
 
 
@@ -36,5 +36,5 @@ class CliHandler():
 
     @staticmethod
     @abstractmethod
-    def handle_cli(options, func):
+    def handle_cli(args, func, options):
         pass
diff --git a/bentoml/handlers/dataframe_handler.py b/bentoml/handlers/dataframe_handler.py
index f8ad0e60..993cba9e 100644
--- a/bentoml/handlers/dataframe_handler.py
+++ b/bentoml/handlers/dataframe_handler.py
@@ -22,10 +22,24 @@ import sys
 import json
 import pandas as pd
 from flask import Response, make_response
-from bentoml.handlers import CliHandler, RequestHandler
+from bentoml.handlers.base_handlers import CliHandler, RequestHandler
+from bentoml.handlers.utils import merge_dicts
 from bentoml.cli.utils import generate_default_parser
 
 
+def check_missing_columns(required_columns, df_columns):
+    missing_columns = set(required_columns) - set(df_columns)
+    if missing_columns:
+        raise ValueError('Missing columns: {}'.format(','.join(missing_columns)))
+
+
+default_options = {
+    'input_json_orient': 'records',
+    'output_json_orient': 'records',
+    'input_columns_require': []
+}
+
+
 class DataframeHandler(RequestHandler, CliHandler):
     """
     Create Data frame handler.  Dataframe handler will take inputs from rest request
@@ -33,18 +47,26 @@ class DataframeHandler(RequestHandler, CliHandler):
     """
 
     @staticmethod
-    def handle_request(request, func):
+    def handle_request(request, func, options={}):
+        options = merge_dicts(default_options, options)
+        if request.headers.get('input_json_orient'):
+            options['input_json_orient'] = request.headers['input_json_orient']
+
         if request.content_type == 'application/json':
-            df = pd.read_json(request.data.decode('utf-8'))
+            df = pd.read_json(
+                request.data.decode('utf-8'), orient=options['input_json_orient'], dtype=False)
         elif request.content_type == 'text/csv':
             df = pd.read_csv(request.data.decode('utf-8'))
         else:
             return make_response(400)
 
+        if options['input_columns_require']:
+            check_missing_columns(options['input_columns_require'], df.columns)
+
         output = func(df)
 
         if isinstance(output, pd.DataFrame):
-            result = output.to_json(orient='records')
+            result = output.to_json(orient=options['output_json_orient'])
         else:
             result = json.dumps(output)
 
@@ -52,21 +74,30 @@ class DataframeHandler(RequestHandler, CliHandler):
         return response
 
     @staticmethod
-    def handle_cli(options, func):
+    def handle_cli(args, func, options={}):
         parser = generate_default_parser()
-        parsed_args = parser.parse_args(options)
+        parser.add_argument('--input_json_orient', default='records')
+        parsed_args = parser.parse_args(args)
+        options = merge_dicts(default_options, options)
+
+        if parsed_args.input_json_orient:
+            options['input_json_orient'] = parsed_args.input_json_orient
 
         with open(parsed_args.input, 'r') as content_file:
             content = content_file.read()
             if content_file.name.endswith('.csv'):
                 df = pd.read_csv(content)
             elif content_file.name.endswith('.json'):
-                df = pd.read_json(content)
+                df = pd.read_json(content, orient=options['input_json_orient'], dtype=False)
+
+            if options['input_columns_require']:
+                check_missing_columns(options['input_columns_require'], df.columns)
+
             output = func(df)
 
             if parsed_args.output == 'json' or not parsed_args.output:
                 if isinstance(output, pd.DataFrame):
-                    result = output.to_json(orient='records')
+                    result = output.to_json(orient=options['output_json_orient'])
                     result = json.loads(result)
                     result = json.dumps(result, indent=2)
                 else:
diff --git a/bentoml/handlers/image_handler.py b/bentoml/handlers/image_handler.py
index 61b15e0c..293e2100 100644
--- a/bentoml/handlers/image_handler.py
+++ b/bentoml/handlers/image_handler.py
@@ -27,9 +27,9 @@ class ImageHandler(RequestHandler, CliHandler):
     """
 
     @staticmethod
-    def handle_request(request, func):
+    def handle_request(request, func, options=None):
         raise NotImplementedError
 
     @staticmethod
-    def handle_cli(options, func):
+    def handle_cli(args, func, options=None):
         raise NotImplementedError
diff --git a/bentoml/handlers/json_handler.py b/bentoml/handlers/json_handler.py
index eace73b5..caee4193 100644
--- a/bentoml/handlers/json_handler.py
+++ b/bentoml/handlers/json_handler.py
@@ -23,8 +23,14 @@ import json
 from flask import Response, make_response
 
 from bentoml.handlers.base_handlers import RequestHandler, CliHandler
+from bentoml.handlers.utils import merge_dicts
 from bentoml.cli.utils import generate_default_parser
 
+default_options = {
+    'input_json_orient': 'records',
+    'output_json_orient': 'records',
+}
+
 
 class JsonHandler(RequestHandler, CliHandler):
     """
@@ -32,7 +38,8 @@ class JsonHandler(RequestHandler, CliHandler):
     """
 
     @staticmethod
-    def handle_request(request, func):
+    def handle_request(request, func, options={}):
+        options = merge_dicts(default_options, options)
         if request.content_type == 'application/json':
             parsed_json = json.loads(request.data.decode('utf-8'))
         else:
@@ -55,9 +62,10 @@ class JsonHandler(RequestHandler, CliHandler):
         return response
 
     @staticmethod
-    def handle_cli(options, func):
+    def handle_cli(args, func, options={}):
+        options = merge_dicts(default_options, options)
         parser = generate_default_parser()
-        parsed_args = parser.parse_args(options)
+        parsed_args = parser.parse_args(args)
 
         with open(parsed_args.input, 'r') as content_file:
             content = content_file.read()
diff --git a/bentoml/handlers/pytorch_tensor_handler.py b/bentoml/handlers/pytorch_tensor_handler.py
index 948f3d3c..7b49485f 100644
--- a/bentoml/handlers/pytorch_tensor_handler.py
+++ b/bentoml/handlers/pytorch_tensor_handler.py
@@ -30,9 +30,9 @@ class PytorchTensorHanlder(RequestHandler, CliHandler):
     """
 
     @staticmethod
-    def handle_request(request, func):
+    def handle_request(request, func, options=None):
         raise NotImplementedError
 
     @staticmethod
-    def handle_cli(options, func):
+    def handle_cli(args, func, options=None):
         raise NotImplementedError
diff --git a/bentoml/handlers/tensorflow_tensor_handler.py b/bentoml/handlers/tensorflow_tensor_handler.py
index 6e81ebd0..44ebef75 100644
--- a/bentoml/handlers/tensorflow_tensor_handler.py
+++ b/bentoml/handlers/tensorflow_tensor_handler.py
@@ -27,9 +27,9 @@ class TensorflowTensorHandler(RequestHandler, CliHandler):
     """
 
     @staticmethod
-    def handle_request(request, func):
+    def handle_request(request, func, options=None):
         raise NotImplementedError
 
     @staticmethod
-    def handle_cli(options, func):
+    def handle_cli(args, func, options=None):
         raise NotImplementedError
diff --git a/bentoml/handlers/utils.py b/bentoml/handlers/utils.py
new file mode 100644
index 00000000..571896ee
--- /dev/null
+++ b/bentoml/handlers/utils.py
@@ -0,0 +1,33 @@
+# BentoML - Machine Learning Toolkit for packaging and deploying models
+# Copyright (C) 2019 Atalaya Tech, Inc.
+
+# This program is free software: you can redistribute it and/or modify
+# it under the terms of the GNU General Public License as published by
+# the Free Software Foundation, either version 3 of the License, or
+# (at your option) any later version.
+
+# This program is distributed in the hope that it will be useful,
+# but WITHOUT ANY WARRANTY; without even the implied warranty of
+# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
+# GNU General Public License for more details.
+
+# You should have received a copy of the GNU General Public License
+# along with this program.  If not, see <http://www.gnu.org/licenses/>.
+
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+import sys
+
+
+def merge_dicts(x, y):
+    """
+    Merge 2 dictionaries into one. The second directory override first dictionary's key
+    """
+    if sys.version_info >= (3, 5):
+        return {**x, **y}
+    else:
+        temp = x.copy()
+        temp.update(y)
+        return temp
diff --git a/bentoml/server/model_api_server.py b/bentoml/server/model_api_server.py
index 26eb6c97..b695912d 100644
--- a/bentoml/server/model_api_server.py
+++ b/bentoml/server/model_api_server.py
@@ -44,14 +44,14 @@ def create_api_function_wrapper(logger, model_name, model_version, api):
     """
     Create api function for flask route
     """
-    summary_name = model_name + '_' + model_version + '_' + api.name
+    summary_name = str(model_name) + '_' + str(model_version) + '_' + str(api.name)
     request_metric_time = Summary(summary_name, summary_name + ' request latency')
 
     def wrapper():
         with request_metric_time.time():
             request_time = time()
             request_id = str(uuid.uuid4())
-            response = api.handler.handle_request(request, api.func)
+            response = api.handler.handle_request(request, api.func, api.options)
             response.headers['request_id'] = request_id
             if response.status_code == 200:
                 metadata = {
diff --git a/bentoml/service.py b/bentoml/service.py
index d818f827..31ed4c61 100644
--- a/bentoml/service.py
+++ b/bentoml/service.py
@@ -59,7 +59,8 @@ class BentoService(object):
                 api_name = _get_func_attr(function, '_api_name')
                 handler = _get_func_attr(function, '_handler')
                 func = function.__get__(self)
-                self._service_apis.append(BentoServiceAPI(api_name, handler, func))
+                options = _get_func_attr(function, '_options')
+                self._service_apis.append(BentoServiceAPI(api_name, handler, func, options))
 
     @property
     @abstractmethod
@@ -77,7 +78,7 @@ class BentoService(object):
         return self._service_apis
 
     @staticmethod
-    def api(handler=DataframeHandler, api_name=None):
+    def api(handler=DataframeHandler, api_name=None, options=None):
         """
         Decorator for adding api to a BentoService
 
@@ -104,6 +105,8 @@ class BentoService(object):
             else:
                 _set_func_attr(func, '_api_name', api_name)
 
+            _set_func_attr(func, '_options', options)
+
             return func
 
         return api_decorator
@@ -115,7 +118,7 @@ class BentoServiceAPI(object):
     with BentoAPIServer and BentoCLI
     """
 
-    def __init__(self, name, handler, func):
+    def __init__(self, name, handler, func, options):
         """
         :param name: API name
         :param handler: A BentoHandler that transforms HTTP Request and/or
@@ -126,6 +129,7 @@ class BentoServiceAPI(object):
         self._name = name
         self._handler = handler
         self._func = func
+        self._options = options
 
     @property
     def name(self):
@@ -138,3 +142,7 @@ class BentoServiceAPI(object):
     @property
     def func(self):
         return self._func
+
+    @property
+    def options(self):
+        return self._options


Test patch:
diff --git a/tests/cli/test_cli.py b/tests/cli/test_cli.py
index 9a0972db..df948f28 100644
--- a/tests/cli/test_cli.py
+++ b/tests/cli/test_cli.py
@@ -5,7 +5,7 @@ sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), '..')
 
 from click.testing import CliRunner
 
-from bentoml.cli import run
+from bentoml.cli.commands import run
 from tests.utils import generate_fake_model
 
 
diff --git a/tests/server/test_model_api_server.py b/tests/server/test_model_api_server.py
index 433930c3..9da6181b 100644
--- a/tests/server/test_model_api_server.py
+++ b/tests/server/test_model_api_server.py
@@ -2,16 +2,20 @@ import os
 import json
 import uuid
 import sys
-sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), '..')))
+sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), '..', '..')))
 
 import bentoml
 from bentoml.server import BentoModelApiServer
 
-
 BASE_TEST_PATH = "/tmp/bentoml-test"
+
+
 class MyFakeModel(object):
-    def predict(self, input):
-        return input.add(3)
+
+    def predict(self, df):
+        df['age'] = df['age'].add(5)
+        return df
+
 
 class RestServiceTestModel(bentoml.BentoModel):
     """
@@ -19,21 +23,19 @@ class RestServiceTestModel(bentoml.BentoModel):
     """
 
     def config(self, artifacts, env):
-        artifacts.add(bentoml.artifacts.PickleArtifact('sentiment_lr'))
+        artifacts.add(bentoml.artifacts.PickleArtifact('fake_model'))
 
-        env.add_conda_dependencies(["scikit-learn"])
-
-    @bentoml.api(bentoml.handlers.DataframeHandler)
+    @bentoml.api(bentoml.handlers.DataframeHandler, options={'input_columns_require': ['age']})
     def predict(self, df):
         """
         predict expects dataframe as input
         """
-        return self.artifacts.sentiment_lr.predict(df)
+        return self.artifacts.fake_model.predict(df)
 
 
 def create_rest_server():
-    sentiment_lr = MyFakeModel()
-    sm = RestServiceTestModel(sentiment_lr=sentiment_lr)
+    fake_model = MyFakeModel()
+    sm = RestServiceTestModel(fake_model=fake_model)
     version = "test_" + uuid.uuid4().hex
     sm.save(BASE_TEST_PATH, version=version)
 
@@ -45,7 +47,6 @@ def create_rest_server():
     return rest_server
 
 
-
 def test_api_function_route():
     rest_server = create_rest_server()
     test_client = rest_server.app.test_client()
@@ -55,11 +56,10 @@ def test_api_function_route():
         list.append(rule.endpoint)
 
     assert 'predict' in list
-    data = [{'age': 12}]
+    data = [{'age': 10}]
 
     response = test_client.post('/predict', data=json.dumps(data), content_type='application/json')
 
-    print(response.data)
     if sys.version_info.major < 3:
         import ast
         loaded_json = json.loads(response.data)
diff --git a/tests/test_model.py b/tests/test_model.py
index 9877d863..abaeaaaa 100644
--- a/tests/test_model.py
+++ b/tests/test_model.py
@@ -63,6 +63,7 @@ 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()