Shell: change server from gunicorn to waitress (#3001)

.github/workflows/style_type_checks.yml

@@ -21,6 +21,7 @@ jobs:
         pip install .
         pip install click black mypy
         pip install types-python-dateutil
+        pip install types-waitress
     - name: Style and type checks
       run: |
         just black

requirements/requirements-extras-shell.txt

@@ -1,2 +1,2 @@
 flask~=2.0
-gunicorn~=19.9
+waitress~=2.1.2

src/gluonts/shell/__main__.py

@@ -17,6 +17,7 @@
 from typing import Optional
 
 import click
+import waitress
 
 from gluonts.env import env as gluonts_env
 from gluonts.shell.serve import Settings
@@ -82,12 +83,13 @@ def serve_command(
     else:
         forecaster_type = None
 
-    gunicorn_app = serve.make_gunicorn_app(
+    flask_app = serve.make_flask_app(
         env=env,
         forecaster_type=forecaster_type,
         settings=Settings(),
     )
-    gunicorn_app.run()
+
+    waitress.serve(flask_app, listen="*:8080")
 
 
 @cli.command(name="train")

src/gluonts/shell/serve/__init__.py

@@ -17,7 +17,6 @@
 from typing import List, Optional, Type, Union
 
 from flask import Flask
-from gunicorn.app.base import BaseApplication
 from pydantic import BaseSettings
 
 import gluonts
@@ -94,32 +93,11 @@ def number_of_workers(self) -> int:
             return cpu_count
 
 
-class Application(BaseApplication):
-    def __init__(self, app, config) -> None:
-        self.app = app
-        self.config = config
-        BaseApplication.__init__(self)
-
-    def load_config(self) -> None:
-        for key, value in self.config.items():
-            if key in self.cfg.settings and value is not None:
-                self.cfg.set(key, value)
-
-    def init(self, parser, opts, args):
-        pass
-
-    def load(self) -> Flask:
-        return self.app
-
-    def stop(self, *args, **kwargs):
-        logger.info("Shutting down GluonTS scoring service")
-
-
-def make_gunicorn_app(
+def make_flask_app(
     env: ServeEnv,
     forecaster_type: Optional[Type[Union[Estimator, Predictor]]],
     settings: Settings,
-) -> Application:
+) -> Flask:
     if forecaster_type is not None:
         logger.info("Using dynamic predictor factory")
 
@@ -159,13 +137,4 @@ def predictor_factory(request) -> Predictor:
         settings=settings,
     )
 
-    gunicorn_app = Application(
-        app=flask_app,
-        config={
-            "bind": settings.sagemaker_server_bind,
-            "workers": settings.number_of_workers,
-            "timeout": settings.sagemaker_server_timeout,
-        },
-    )
-
-    return gunicorn_app
+    return flask_app

src/gluonts/testutil/shell.py

@@ -19,6 +19,7 @@
 import tempfile
 import time
 import typing
+import waitress
 from contextlib import closing, contextmanager
 from dataclasses import dataclass
 from multiprocessing.context import ForkContext
@@ -33,7 +34,7 @@
 from gluonts.shell.env import ServeEnv, TrainEnv
 from gluonts.shell.sagemaker import ServePaths, TrainPaths
 from gluonts.shell.sagemaker.params import encode_sagemaker_parameters
-from gluonts.shell.serve import Settings, make_gunicorn_app
+from gluonts.shell.serve import Settings, make_flask_app
 
 
 class ServerFacade:
@@ -121,10 +122,12 @@ class Server:
     settings: Settings = Settings()
 
     def run(self):
-        gunicorn_app = make_gunicorn_app(
+        flask_app = make_flask_app(
             self.env, self.forecaster_type, self.settings
         )
-        gunicorn_app.run()
+        waitress.serve(
+            flask_app, listen=f"*:{self.settings.sagemaker_server_port}"
+        )
 
 
 @contextmanager

test/shell/require-packages.txt

@@ -1,3 +1,3 @@
 flask
-gunicorn
+waitress
 requests