feat: Add gunicorn for serve with multiprocess (feast-dev#3636)

* feat: Add gunicorn for serve with multiprocess Signed-off-by: phil.park <[email protected]> * feat: Add gunicorn to requirements for ci Signed-off-by: phil.park <[email protected]> --------- Signed-off-by: phil.park <[email protected]>
zseta · Jun 5, 2023 · 4de7faf · 4de7faf
1 parent bf740d2
commit 4de7faf
Show file tree

Hide file tree

Showing 10 changed files with 81 additions and 10 deletions.
diff --git a/sdk/python/feast/cli.py b/sdk/python/feast/cli.py
@@ -649,6 +649,21 @@ def init_command(project_directory, minimal: bool, template: str):
     show_default=True,
     help="Disable logging served features",
 )
+@click.option(
+    "--workers",
+    "-w",
+    type=click.INT,
+    default=1,
+    show_default=True,
+    help="Number of worker",
+)
+@click.option(
+    "--keep-alive-timeout",
+    type=click.INT,
+    default=5,
+    show_default=True,
+    help="Timeout for keep alive",
+)
 @click.pass_context
 def serve_command(
     ctx: click.Context,
@@ -657,11 +672,21 @@ def serve_command(
     type_: str,
     no_access_log: bool,
     no_feature_log: bool,
+    workers: int,
+    keep_alive_timeout: int,
 ):
     """Start a feature server locally on a given port."""
     store = create_feature_store(ctx)
 
-    store.serve(host, port, type_, no_access_log, no_feature_log)
+    store.serve(
+        host=host,
+        port=port,
+        type_=type_,
+        no_access_log=no_access_log,
+        no_feature_log=no_feature_log,
+        workers=workers,
+        keep_alive_timeout=keep_alive_timeout,
+    )
 
 
 @cli.command("serve_transformations")

diff --git a/sdk/python/feast/feature_server.py b/sdk/python/feast/feature_server.py
@@ -2,8 +2,8 @@
 import traceback
 import warnings
 
+import gunicorn.app.base
 import pandas as pd
-import uvicorn
 from fastapi import FastAPI, HTTPException, Request, Response, status
 from fastapi.logger import logger
 from fastapi.params import Depends
@@ -137,8 +137,35 @@ def health():
     return app
 
 
+class FeastServeApplication(gunicorn.app.base.BaseApplication):
+    def __init__(self, store: "feast.FeatureStore", **options):
+        self._app = get_app(store=store)
+        self._options = options
+        super().__init__()
+
+    def load_config(self):
+        for key, value in self._options.items():
+            if key.lower() in self.cfg.settings and value is not None:
+                self.cfg.set(key.lower(), value)
+
+        self.cfg.set("worker_class", "uvicorn.workers.UvicornWorker")
+
+    def load(self):
+        return self._app
+
+
 def start_server(
-    store: "feast.FeatureStore", host: str, port: int, no_access_log: bool
+    store: "feast.FeatureStore",
+    host: str,
+    port: int,
+    no_access_log: bool,
+    workers: int,
+    keep_alive_timeout: int,
 ):
-    app = get_app(store)
-    uvicorn.run(app, host=host, port=port, access_log=(not no_access_log))
+    FeastServeApplication(
+        store=store,
+        bind=f"{host}:{port}",
+        accesslog=None if no_access_log else "-",
+        workers=workers,
+        keepalive=keep_alive_timeout,
+    ).run()
diff --git a/sdk/python/feast/feature_store.py b/sdk/python/feast/feature_store.py
@@ -2152,7 +2152,6 @@ def _get_feature_views_to_use(
         allow_cache=False,
         hide_dummy_entity: bool = True,
     ) -> Tuple[List[FeatureView], List[RequestFeatureView], List[OnDemandFeatureView]]:
-
         fvs = {
             fv.name: fv
             for fv in [
@@ -2223,6 +2222,8 @@ def serve(
         type_: str,
         no_access_log: bool,
         no_feature_log: bool,
+        workers: int,
+        keep_alive_timeout: int,
     ) -> None:
         """Start the feature consumption server locally on a given port."""
         type_ = type_.lower()
@@ -2231,7 +2232,14 @@ def serve(
                 f"Python server only supports 'http'. Got '{type_}' instead."
             )
         # Start the python server
-        feature_server.start_server(self, host, port, no_access_log)
+        feature_server.start_server(
+            self,
+            host=host,
+            port=port,
+            no_access_log=no_access_log,
+            workers=workers,
+            keep_alive_timeout=keep_alive_timeout,
+        )
 
     @log_exceptions_and_usage
     def get_feature_server_endpoint(self) -> Optional[str]:

diff --git a/sdk/python/requirements/py3.10-ci-requirements.txt b/sdk/python/requirements/py3.10-ci-requirements.txt
@@ -326,6 +326,8 @@ grpcio-testing==1.54.2
     # via feast (setup.py)
 grpcio-tools==1.54.2
     # via feast (setup.py)
+gunicorn==20.1.0
+    # via feast (setup.py)
 h11==0.14.0
     # via
     #   httpcore

diff --git a/sdk/python/requirements/py3.10-requirements.txt b/sdk/python/requirements/py3.10-requirements.txt
@@ -59,6 +59,8 @@ grpcio==1.54.2
     #   grpcio-reflection
 grpcio-reflection==1.54.2
     # via feast (setup.py)
+gunicorn==20.1.0
+    # via feast (setup.py)
 h11==0.14.0
     # via
     #   httpcore

diff --git a/sdk/python/requirements/py3.8-ci-requirements.txt b/sdk/python/requirements/py3.8-ci-requirements.txt
@@ -330,6 +330,8 @@ grpcio-testing==1.54.2
     # via feast (setup.py)
 grpcio-tools==1.54.2
     # via feast (setup.py)
+gunicorn==20.1.0
+    # via feast (setup.py)
 h11==0.14.0
     # via
     #   httpcore

diff --git a/sdk/python/requirements/py3.8-requirements.txt b/sdk/python/requirements/py3.8-requirements.txt
@@ -59,6 +59,8 @@ grpcio==1.54.2
     #   grpcio-reflection
 grpcio-reflection==1.54.2
     # via feast (setup.py)
+gunicorn==20.1.0
+    # via feast (setup.py)
 h11==0.14.0
     # via
     #   httpcore

diff --git a/sdk/python/requirements/py3.9-ci-requirements.txt b/sdk/python/requirements/py3.9-ci-requirements.txt
@@ -326,6 +326,8 @@ grpcio-testing==1.54.2
     # via feast (setup.py)
 grpcio-tools==1.54.2
     # via feast (setup.py)
+gunicorn==20.1.0
+    # via feast (setup.py)
 h11==0.14.0
     # via
     #   httpcore

diff --git a/sdk/python/requirements/py3.9-requirements.txt b/sdk/python/requirements/py3.9-requirements.txt
@@ -59,6 +59,8 @@ grpcio==1.54.2
     #   grpcio-reflection
 grpcio-reflection==1.54.2
     # via feast (setup.py)
+gunicorn==20.1.0
+    # via feast (setup.py)
 h11==0.14.0
     # via
     #   httpcore

diff --git a/setup.py b/setup.py
@@ -70,6 +70,7 @@
     "typeguard==2.13.3",
     "fastapi>=0.68.0,<1",
     "uvicorn[standard]>=0.14.0,<1",
+    "gunicorn",
     "dask>=2021.1.0",
     "bowler",  # Needed for automatic repo upgrades
     # FastAPI does not correctly pull starlette dependency on httpx see thread(https://github.com/tiangolo/fastapi/issues/5656).
@@ -103,9 +104,7 @@
     "pyspark>=3.0.0,<4",
 ]
 
-TRINO_REQUIRED = [
-    "trino>=0.305.0,<0.400.0", "regex"
-]
+TRINO_REQUIRED = ["trino>=0.305.0,<0.400.0", "regex"]
 
 POSTGRES_REQUIRED = [
     "psycopg2-binary>=2.8.3,<3",