initialize model once

src/sagemaker_chainer_container/serving.py

@@ -105,17 +105,18 @@ def _user_module_transformer(user_module):
                                    output_fn=output_fn)
 
 
-def main(environ, start_response):
-    serving_env = env.ServingEnv()
-
-    logger.setLevel(serving_env.log_level)
+app = None
 
-    user_module = modules.import_module(serving_env.module_dir, serving_env.module_name)
 
-    user_module_transformer = _user_module_transformer(user_module)
-
-    user_module_transformer.initialize()
+def main(environ, start_response):
+    global app
+    if app is None:
+        serving_env = env.ServingEnv()
+        logger.setLevel(serving_env.log_level)
+        user_module = modules.import_module(serving_env.module_dir, serving_env.module_name)
+        user_module_transformer = _user_module_transformer(user_module)
+        user_module_transformer.initialize()
+        app = worker.Worker(transform_fn=user_module_transformer.transform,
+                            module_name=serving_env.module_name)
 
-    app = worker.Worker(transform_fn=user_module_transformer.transform,
-                        module_name=serving_env.module_name)
     return app(environ, start_response)

test/integration/local/test_serving.py

@@ -0,0 +1,21 @@
+from __future__ import absolute_import
+
+import os
+
+import requests
+
+from test.utils import local_mode
+
+path = os.path.dirname(os.path.realpath(__file__))
+resources_path = os.path.abspath(os.path.join(path, '..', '..', 'resources'))
+
+
+def test_serving_calls_model_fn_once(docker_image, opt_ml):
+    script_path = os.path.join(resources_path, 'call_model_fn_once.py')
+    with local_mode.serve(script_path, model_dir=None, image_name=docker_image, opt_ml=opt_ml,
+                          additional_env_vars=['SAGEMAKER_MODEL_SERVER_WORKERS=2']):
+
+        # call enough times to ensure multiple requests to a worker
+        for i in range(3):
+            # will return 500 error if model_fn called during request handling
+            assert b'output' == requests.post(local_mode.REQUEST_URL, data=b'input').content

test/resources/call_model_fn_once.py

@@ -0,0 +1,36 @@
+# Copyright 2018 Amazon.com, Inc. or its affiliates. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License"). You
+# may not use this file except in compliance with the License. A copy of
+# the License is located at
+#
+#     http://aws.amazon.com/apache2.0/
+#
+# or in the "license" file accompanying this file. This file 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.
+
+import os
+
+
+def model_fn(model_dir):
+    lock_file = os.path.join(model_dir, 'model_fn.lock.{}'.format(os.getpid()))
+    if os.path.exists(lock_file):
+        raise RuntimeError('model_fn called more than once (lock: {})'.format(lock_file))
+
+    open(lock_file, 'a').close()
+
+    return 'model'
+
+
+def input_fn(data, content_type):
+    return data
+
+
+def predict_fn(data, model):
+    return b'output'
+
+
+def output_fn(prediction, accept):
+    return prediction