From dc75f71be476d43d1f86a8956f76a05f9b384bd7 Mon Sep 17 00:00:00 2001 From: Marcio Dos Santos Date: Thu, 3 May 2018 09:32:14 -0700 Subject: [PATCH 1/2] Add worker --- setup.py | 2 +- src/sagemaker_containers/__init__.py | 11 ++- src/sagemaker_containers/content_types.py | 13 +++ src/sagemaker_containers/status_codes.py | 14 ++++ src/sagemaker_containers/transformer.py | 22 +++++ src/sagemaker_containers/worker.py | 80 +++++++++++++++++++ test/functional/test_worker_with_transform.py | 64 +++++++++++++++ test/unit/test_environment.py | 4 +- test/unit/test_server.py | 1 - test/unit/test_worker.py | 66 +++++++++++++++ 10 files changed, 271 insertions(+), 6 deletions(-) create mode 100644 src/sagemaker_containers/content_types.py create mode 100644 src/sagemaker_containers/status_codes.py create mode 100644 src/sagemaker_containers/transformer.py create mode 100644 src/sagemaker_containers/worker.py create mode 100644 test/functional/test_worker_with_transform.py create mode 100644 test/unit/test_worker.py diff --git a/setup.py b/setup.py index e50df45..d44e724 100644 --- a/setup.py +++ b/setup.py @@ -37,7 +37,7 @@ def read(file_name): 'Programming Language :: Python :: 2.7', 'Programming Language :: Python :: 3.5', ], - install_requires=['boto3', 'six', 'pip'], + install_requires=['boto3', 'six', 'pip', 'flask', 'gunicorn', 'gevent'], extras_require={ 'test': ['tox', 'flake8', 'pytest', 'pytest-cov', 'mock', 'sagemaker', 'numpy'] diff --git a/src/sagemaker_containers/__init__.py b/src/sagemaker_containers/__init__.py index 45d4d62..9e37a2d 100644 --- a/src/sagemaker_containers/__init__.py +++ b/src/sagemaker_containers/__init__.py @@ -12,7 +12,14 @@ # language governing permissions and limitations under the License. from __future__ import absolute_import -from sagemaker_containers import collections, environment, functions, modules, server # noqa ignore=F401 -# imported but unused +import sagemaker_containers.collections +import sagemaker_containers.content_types +import sagemaker_containers.environment +import sagemaker_containers.functions +import sagemaker_containers.modules +import sagemaker_containers.server +import sagemaker_containers.status_codes +import sagemaker_containers.worker # noqa ignore=F401 + from sagemaker_containers.environment import Environment, ServingEnvironment, TrainingEnvironment # noqa ignore=F401 # imported but unused diff --git a/src/sagemaker_containers/content_types.py b/src/sagemaker_containers/content_types.py new file mode 100644 index 0000000..3d6fb40 --- /dev/null +++ b/src/sagemaker_containers/content_types.py @@ -0,0 +1,13 @@ +# 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. +APPLICATION_JSON = 'application/json' diff --git a/src/sagemaker_containers/status_codes.py b/src/sagemaker_containers/status_codes.py new file mode 100644 index 0000000..dbcb129 --- /dev/null +++ b/src/sagemaker_containers/status_codes.py @@ -0,0 +1,14 @@ +# 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. +OK = 200 +ACCEPTED = 202 diff --git a/src/sagemaker_containers/transformer.py b/src/sagemaker_containers/transformer.py new file mode 100644 index 0000000..33ebd10 --- /dev/null +++ b/src/sagemaker_containers/transformer.py @@ -0,0 +1,22 @@ +# 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. +from __future__ import absolute_import + + +class Transformer(object): + + def initialize(self): + pass + + def transform(self): + pass diff --git a/src/sagemaker_containers/worker.py b/src/sagemaker_containers/worker.py new file mode 100644 index 0000000..5d97cbf --- /dev/null +++ b/src/sagemaker_containers/worker.py @@ -0,0 +1,80 @@ +# 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. +from __future__ import absolute_import + +import collections + +from flask import Flask, Response + +import sagemaker_containers as smc + +env = smc.environment.ServingEnvironment() + + +def default_healthcheck_fn(): # type: () -> Response + """Ping is default health-check handler. Returns 200 with no content. + + During a new serving container startup, Amazon SageMaker starts sending periodic GET requests to the /ping endpoint + to ensure that the container is ready for predictions. + + The simplest requirement on the container is to respond with an HTTP 200 status code and an empty body. This + indicates to Amazon SageMaker that the container is ready to accept inference requests at the /invocations endpoint. + + If the container does not begin to consistently respond with 200s during the first 30 seconds after startup, + the CreateEndPoint and UpdateEndpoint APIs will fail. + + While the minimum bar is for the container to return a static 200, a container developer can use this functionality + to perform deeper checks. The request timeout on /ping attempts is 2 seconds. + + More information on how health-check works can be found here: + https://docs.aws.amazon.com/sagemaker/latest/dg/your-algorithms-inference-code.html#your-algorithms-inference-algo-ping-requests + + Returns: + (flask.Response): with status code 200 + """ + return Response(status=smc.status_codes.OK) + + +def run(transformer, healthcheck_fn=None, module_name=None): + # type: (smc.Transformer, function or None, str or None) -> Flask + """Creates and Flask application from a transformer. + + Args: + transformer (smc.Transformer): object responsible to load the model and make predictions. + healthcheck_fn (function): function that will be used for healthcheck calls when the containers starts, + if not specified, it will use ping as the default healthcheck call. + module_name (str): the module name which implements the worker. If not specified, ir will use + sagemaker_containers.ServingEnvironment().module_name as the default module name. + + Returns: + (Flask): an instance of Flask ready for inferences. + """ + healthcheck_fn = healthcheck_fn or default_healthcheck_fn + app = Flask(import_name=module_name or env.module_name) + + transformer.initialize() + + def invocations_fn(): + transform_spec = transformer.transform() + + return Response(response=transform_spec.prediction, + status=smc.status_codes.OK, + mimetype=transform_spec.accept) + + app.add_url_rule(rule='/invocations', endpoint='invocations', view_func=invocations_fn, methods=["POST"]) + app.add_url_rule(rule='/ping', endpoint='ping', view_func=healthcheck_fn or default_healthcheck_fn) + + return app + + +TransformSpec = collections.namedtuple('TransformSpec', 'prediction accept') diff --git a/test/functional/test_worker_with_transform.py b/test/functional/test_worker_with_transform.py new file mode 100644 index 0000000..d76870d --- /dev/null +++ b/test/functional/test_worker_with_transform.py @@ -0,0 +1,64 @@ +# 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. +from __future__ import absolute_import + +import json + +from mock import patch, PropertyMock +import pytest +from six.moves import range + +import sagemaker_containers as smc + + +class DummyTransformer(object): + def __init__(self): + self.calls = dict(initialize=0, transform=0) + + def initialize(self): + self.calls['initialize'] += 1 + + def transform(self): + self.calls['transform'] += 1 + return smc.worker.TransformSpec(json.dumps(self.calls), smc.content_types.APPLICATION_JSON) + + +@patch('sagemaker_containers.environment.ServingEnvironment.module_name', PropertyMock(return_value='user_program')) +@pytest.mark.parametrize('module_name,expected_name', [('my_module', 'my_module'), (None, 'user_program')]) +def test_worker(module_name, expected_name): + transformer = DummyTransformer() + + with smc.worker.run(transformer, module_name=module_name).test_client() as worker: + assert worker.application.import_name == expected_name + + assert worker.get('/ping').status_code == smc.status_codes.OK + + for _ in range(9): + response = worker.post('/invocations') + assert response.status_code == smc.status_codes.OK + + response = worker.post('/invocations') + assert json.loads(response.get_data().decode('utf-8')) == dict(initialize=1, transform=10) + assert response.mimetype == smc.content_types.APPLICATION_JSON + + +def test_worker_with_custom_ping(): + transformer = DummyTransformer() + + def custom_ping(): + return 'ping', smc.status_codes.ACCEPTED + + with smc.worker.run(transformer, custom_ping, 'custom_ping').test_client() as worker: + response = worker.get('/ping') + assert response.status_code == smc.status_codes.ACCEPTED + assert response.get_data().decode('utf-8') == 'ping' diff --git a/test/unit/test_environment.py b/test/unit/test_environment.py index 2ac8e50..7843183 100644 --- a/test/unit/test_environment.py +++ b/test/unit/test_environment.py @@ -221,7 +221,7 @@ def test_environment_module_name(sagemaker_program): @patch('tempfile.mkdtemp') @patch('shutil.rmtree') -def test_temporary_directory(rmtree, mkdtemp): +def test_tmpdir(rmtree, mkdtemp): with smc.environment.tmpdir(): mkdtemp.assert_called() rmtree.assert_called() @@ -229,7 +229,7 @@ def test_temporary_directory(rmtree, mkdtemp): @patch('tempfile.mkdtemp') @patch('shutil.rmtree') -def test_temporary_directory_with_args(rmtree, mkdtemp): +def test_tmpdir_with_args(rmtree, mkdtemp): with smc.environment.tmpdir('suffix', 'prefix', '/tmp'): mkdtemp.assert_called_with(dir='/tmp', prefix='prefix', suffix='suffix') rmtree.assert_called() diff --git a/test/unit/test_server.py b/test/unit/test_server.py index 718dad2..7d1aa5e 100644 --- a/test/unit/test_server.py +++ b/test/unit/test_server.py @@ -15,7 +15,6 @@ from sagemaker_containers.environment import ServingEnvironment from sagemaker_containers.server import start - @patch.object(ServingEnvironment, 'model_server_workers', PropertyMock(return_value=2)) @patch.object(ServingEnvironment, 'model_server_timeout', PropertyMock(return_value=100)) @patch.object(ServingEnvironment, 'use_nginx', PropertyMock(return_value=False)) diff --git a/test/unit/test_worker.py b/test/unit/test_worker.py new file mode 100644 index 0000000..701ff72 --- /dev/null +++ b/test/unit/test_worker.py @@ -0,0 +1,66 @@ +# 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. +from __future__ import absolute_import + +from mock import ANY, Mock, patch, PropertyMock +import pytest +from six.moves import range + +import sagemaker_containers as smc + + +class MockTransformer(Mock): + def transform(self): + return smc.worker.TransformSpec('fake data', smc.content_types.APPLICATION_JSON) + + +def test_default_ping_fn(): + assert smc.worker.default_healthcheck_fn().status_code == smc.status_codes.OK + + +@pytest.fixture(name='flask') +def patch_flask(): + property_mock = PropertyMock(return_value='user_program') + with patch('sagemaker_containers.worker.Flask') as flask, \ + patch('sagemaker_containers.environment.ServingEnvironment.module_name', + property_mock): + yield flask + + +@pytest.mark.parametrize('module_name, expected_name', [('test_module', 'test_module'), (None, 'user_program')]) +def test_run(flask, module_name, expected_name): + transformer = MockTransformer() + app = smc.worker.run(transformer, module_name=module_name) + + flask.assert_called_with(import_name=expected_name) + + transformer.initialize.assert_called() + + rules = app.add_url_rule + rules.assert_any_call(rule='/invocations', endpoint='invocations', view_func=ANY, methods=ANY) + + rules.assert_called_with(rule='/ping', endpoint='ping', view_func=smc.worker.default_healthcheck_fn) + + assert rules.call_count == 2 + + +def test_invocations(): + transformer = MockTransformer() + app = smc.worker.run(transformer, module_name='test_module') + + with app.test_client() as worker: + for _ in range(9): + response = worker.post('/invocations') + assert response.status_code == smc.status_codes.OK + assert response.get_data().decode('utf-8') == 'fake data' + assert response.mimetype == smc.content_types.APPLICATION_JSON From c3c370d1f45e78dc8abe3167a556485f4bddcbbc Mon Sep 17 00:00:00 2001 From: Marcio Dos Santos Date: Thu, 3 May 2018 09:32:50 -0700 Subject: [PATCH 2/2] WIP --- old | 71 ++++++++++++++++++++ setup.py | 2 +- src/sagemaker_containers/server.py | 68 ++++++++++--------- test/functional/simple_flask.py | 22 +++---- test/functional/test_server.py | 42 +++++------- test/functional/test_server_with_worker.py | 76 ++++++++++++++++++++++ tox.ini | 1 + 7 files changed, 208 insertions(+), 74 deletions(-) create mode 100644 old create mode 100644 test/functional/test_server_with_worker.py diff --git a/old b/old new file mode 100644 index 0000000..044ccbc --- /dev/null +++ b/old @@ -0,0 +1,71 @@ +# 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. +from __future__ import absolute_import + +import subprocess + +import gunicorn.app.base +import pkg_resources + +import sagemaker_containers as smc + +UNIX_SOCKET_BIND = 'unix:/tmp/gunicorn.sock' +HTTP_BIND = '0.0.0.0:8080' + + +class GunicornApp(gunicorn.app.base.BaseApplication): + """Standalone gunicorn application + + """ + + def __init__(self, app, timeout=None, worker_class=None, bind=None, + worker_connections=None, workers=None, log_level=None): + self.log_level = log_level + self.workers = workers + self.worker_connections = worker_connections + self.bind = bind + self.worker_class = worker_class + self.timeout = timeout + self.application = app + super(GunicornApp, self).__init__() + + def load_config(self): + self.cfg.set('timeout', self.timeout) + self.cfg.set('worker_class', self.worker_class) + self.cfg.set('bind', self.bind) + self.cfg.set('worker_connections', self.worker_connections) + self.cfg.set('workers', self.workers) + self.cfg.set('loglevel', self.log_level) + + def load(self): + return self.application + + +def start(app): + env = smc.environment.ServingEnvironment() + gunicorn_bind_address = HTTP_BIND + processes = [] + + if env.use_nginx: + gunicorn_bind_address = UNIX_SOCKET_BIND + nginx_config_file = pkg_resources.resource_filename(smc.__name__, '/etc/nginx.conf') + nginx = subprocess.Popen(['nginx', '-c', nginx_config_file]) + processes.append(nginx) + + try: + GunicornApp(app=app, timeout=env.model_server_timeout, worker_class='gevent', + bind=gunicorn_bind_address, worker_connections=1000 * env.model_server_workers, + workers=env.model_server_workers, log_level='debug').run() + finally: + for p in processes: + p.terminate() diff --git a/setup.py b/setup.py index d44e724..20d1e11 100644 --- a/setup.py +++ b/setup.py @@ -40,7 +40,7 @@ def read(file_name): install_requires=['boto3', 'six', 'pip', 'flask', 'gunicorn', 'gevent'], extras_require={ - 'test': ['tox', 'flake8', 'pytest', 'pytest-cov', 'mock', 'sagemaker', 'numpy'] + 'test': ['tox', 'flake8', 'pytest', 'pytest-cov', 'mock', 'sagemaker', 'numpy', 'requests'] }, entry_points={ diff --git a/src/sagemaker_containers/server.py b/src/sagemaker_containers/server.py index 7c8d0d8..044ccbc 100644 --- a/src/sagemaker_containers/server.py +++ b/src/sagemaker_containers/server.py @@ -12,10 +12,9 @@ # language governing permissions and limitations under the License. from __future__ import absolute_import -import signal import subprocess -import sys +import gunicorn.app.base import pkg_resources import sagemaker_containers as smc @@ -24,44 +23,49 @@ HTTP_BIND = '0.0.0.0:8080' -def add_terminate_signal(process): - def terminate(signal_number, stack_frame): - process.terminate() +class GunicornApp(gunicorn.app.base.BaseApplication): + """Standalone gunicorn application - signal.signal(signal.SIGTERM, terminate) + """ + def __init__(self, app, timeout=None, worker_class=None, bind=None, + worker_connections=None, workers=None, log_level=None): + self.log_level = log_level + self.workers = workers + self.worker_connections = worker_connections + self.bind = bind + self.worker_class = worker_class + self.timeout = timeout + self.application = app + super(GunicornApp, self).__init__() -def start(module_app): + def load_config(self): + self.cfg.set('timeout', self.timeout) + self.cfg.set('worker_class', self.worker_class) + self.cfg.set('bind', self.bind) + self.cfg.set('worker_connections', self.worker_connections) + self.cfg.set('workers', self.workers) + self.cfg.set('loglevel', self.log_level) + def load(self): + return self.application + + +def start(app): env = smc.environment.ServingEnvironment() gunicorn_bind_address = HTTP_BIND - - nginx = None + processes = [] if env.use_nginx: gunicorn_bind_address = UNIX_SOCKET_BIND nginx_config_file = pkg_resources.resource_filename(smc.__name__, '/etc/nginx.conf') nginx = subprocess.Popen(['nginx', '-c', nginx_config_file]) - - add_terminate_signal(nginx) - - gunicorn = subprocess.Popen(['gunicorn', - '--timeout', str(env.model_server_timeout), - '-k', 'gevent', - '-b', gunicorn_bind_address, - '--worker-connections', str(1000 * env.model_server_workers), - '-w', str(env.model_server_workers), - '--log-level', 'info', - module_app]) - - add_terminate_signal(gunicorn) - - while True: - if nginx and nginx.poll(): - nginx.terminate() - break - elif gunicorn.poll(): - gunicorn.terminate() - break - - sys.exit(0) + processes.append(nginx) + + try: + GunicornApp(app=app, timeout=env.model_server_timeout, worker_class='gevent', + bind=gunicorn_bind_address, worker_connections=1000 * env.model_server_workers, + workers=env.model_server_workers, log_level='debug').run() + finally: + for p in processes: + p.terminate() diff --git a/test/functional/simple_flask.py b/test/functional/simple_flask.py index 309a198..50f5464 100644 --- a/test/functional/simple_flask.py +++ b/test/functional/simple_flask.py @@ -10,20 +10,11 @@ # 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 - from flask import Flask -app = Flask(__name__) - - -@app.route('/') -def hello(): - return 'Hello World!' +import sagemaker_containers as smc -@app.route('/invocations') -def invocations(): - return 'invocation' +app = Flask(__name__) @app.route('/ping') @@ -31,6 +22,9 @@ def ping(): return ':)' -@app.route('/shutdown') -def shutdown(): - os._exit(4) +def start_server(): + smc.server.start(app) + + +if __name__ == '__main__': + start_server() diff --git a/test/functional/test_server.py b/test/functional/test_server.py index ae325d1..269dbdc 100644 --- a/test/functional/test_server.py +++ b/test/functional/test_server.py @@ -11,40 +11,28 @@ # ANY KIND, either express or implied. See the License for the specific # language governing permissions and limitations under the License. import os -import threading +import subprocess +import sys import time -import urllib3 +import requests import sagemaker_containers as smc -def test_server(): - os.environ[smc.environment.FRAMEWORK_MODULE_ENV] = 'test.functional.simple_flask:app' - os.environ[smc.environment.USE_NGINX_ENV] = 'false' - - env = smc.environment.ServingEnvironment() - - def worker(): - smc.server.start(env.framework_module) - - t = threading.Thread(target=worker) - t.start() +CURRENT_DIR = os.path.join(os.path.dirname(__file__)) - time.sleep(2) - http = urllib3.PoolManager() - base_url = 'http://127.0.0.1:8080' - r = http.request('GET', '{}/ping'.format(base_url)) - assert r.status == 200 - - r = http.request('GET', '{}/invocations'.format(base_url)) - assert r.status == 200 - assert r.data.decode('utf-8') == 'invocation' +def test_server(): + environ = os.environ.copy() + environ[smc.environment.USE_NGINX_ENV] = 'false' - # shut down the server or else it will go on forever. + application_path = os.path.join(CURRENT_DIR, 'simple_flask.py') + process = subprocess.Popen(args=[sys.executable, application_path], env=environ) try: - http.request('GET', '{}/shutdown'.format(base_url)) - except urllib3.exceptions.MaxRetryError: - # the above request will kill the server so it is expected that it fails. - pass + time.sleep(2) + + assert requests.get('http://127.0.0.1:8080/ping').status_code == 200 + finally: + process.terminate() + time.sleep(2) diff --git a/test/functional/test_server_with_worker.py b/test/functional/test_server_with_worker.py new file mode 100644 index 0000000..0f1a3ff --- /dev/null +++ b/test/functional/test_server_with_worker.py @@ -0,0 +1,76 @@ +# 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 json +import os +import threading +import time + +from mock import patch +import urllib3 + +import sagemaker_containers as smc + + +class DummyTransformer(object): + def __init__(self): + self.calls = dict(initialize=0, transform=0) + + def initialize(self): + self.calls['initialize'] += 1 + + def transform(self): + self.calls['transform'] += 1 + return smc.worker.TransformSpec(json.dumps(self.calls), smc.content_types.APPLICATION_JSON) + + +app = smc.worker.run(DummyTransformer(), module_name='app') + + +# class TestThread(threading.Thread): +# +# def __init__(self, name='TestThread'): +# super(self, TestThread).__init__(name) +# self.event = threading.Event() +# +# def run(self): +# smc.server.start('test.functional.test_server_with_worker:app') +# +# def join(self, timeout=None): +# """ Stop the thread. """ +# self._stopevent.set() +# threading.Thread.join(self, timeout) + + +@patch.dict(os.environ, {smc.environment.FRAMEWORK_MODULE_ENV: 'test.functional.test_server_with_worker:app', + smc.environment.USE_NGINX_ENV: 'false'}) +def test_server(): + def start_server(stop_event): + smc.server.start('test.functional.test_server_with_worker:app') + + pill2kill = threading.Event() + t = threading.Thread(target=start_server, args=(pill2kill)) + t.start() + + time.sleep(2) + + http = urllib3.PoolManager() + base_url = 'http://127.0.0.1:8080' + r = http.request('GET', '{}/ping'.format(base_url)) + assert r.status == 200 + + r = http.request('POST', '{}/invocations'.format(base_url)) + assert r.status == 200 + assert r.data.decode('utf-8') == dict(initialize=1, transform=1) + + pill2kill.set() + t.join() diff --git a/tox.ini b/tox.ini index ba7ae81..7faac4e 100644 --- a/tox.ini +++ b/tox.ini @@ -44,6 +44,7 @@ deps = flask gunicorn gevent + requests [testenv:flake8] basepython = python