Repository: jina-ai/serve
Pull request: feat: add FastAPIBaseGateway
Create a FastAPIBaseGateway where the user just needs to implement a fast API app.
All the rest (setting up, running and shutting down the server) is automatically handled

Also make sure the HTTPGateway inherits from it instead of the BaseGateway

closes: https://github.com/jina-ai/jina/issues/5454

Problem statement:
create FastAPIBaseGateway
Create a FastAPIBaseGateway where the user just needs to implement a fast API app.
All the rest (setting up, running and shutting down the server) is automatically handled

Also make sure the HTTPGateway inherits from it instead of the BaseGateway

Hints from issue discussion:

Commit message: test: fix fastapi gateway test

Patch to review:
diff --git a/jina/serve/runtimes/gateway/http/fastapi.py b/jina/serve/runtimes/gateway/http/fastapi.py
new file mode 100644
index 0000000000000..67ef20261232f
--- /dev/null
+++ b/jina/serve/runtimes/gateway/http/fastapi.py
@@ -0,0 +1,115 @@
+import logging
+import os
+from abc import abstractmethod
+from typing import Optional
+
+from jina.importer import ImportExtensions
+from jina.serve.gateway import BaseGateway
+
+
+class FastAPIBaseGateway(BaseGateway):
+    """HTTP Gateway implementation"""
+
+    def __init__(
+        self,
+        ssl_keyfile: Optional[str] = None,
+        ssl_certfile: Optional[str] = None,
+        uvicorn_kwargs: Optional[dict] = None,
+        proxy: Optional[bool] = None,
+        **kwargs
+    ):
+        """Initialize the FastAPIBaseGateway
+        :param ssl_keyfile: the path to the key file
+        :param ssl_certfile: the path to the certificate file
+        :param uvicorn_kwargs: Dictionary of kwargs arguments that will be passed to Uvicorn server when starting the server
+        :param proxy: If set, respect the http_proxy and https_proxy environment variables, otherwise, it will unset
+            these proxy variables before start. gRPC seems to prefer no proxy
+        :param kwargs: keyword args
+        """
+        super().__init__(**kwargs)
+        self.ssl_keyfile = ssl_keyfile
+        self.ssl_certfile = ssl_certfile
+        self.uvicorn_kwargs = uvicorn_kwargs
+
+        if not proxy and os.name != 'nt':
+            os.unsetenv('http_proxy')
+            os.unsetenv('https_proxy')
+
+    @property
+    @abstractmethod
+    def app(self):
+        '''Get a FastAPI app'''
+        ...
+
+    async def setup_server(self):
+        """
+        Initialize and return GRPC server
+        """
+        with ImportExtensions(required=True):
+            from uvicorn import Config, Server
+
+        class UviServer(Server):
+            """The uvicorn server."""
+
+            async def setup(self, sockets=None):
+                """
+                Setup uvicorn server.
+
+                :param sockets: sockets of server.
+                """
+                config = self.config
+                if not config.loaded:
+                    config.load()
+                self.lifespan = config.lifespan_class(config)
+                await self.startup(sockets=sockets)
+                if self.should_exit:
+                    return
+
+            async def serve(self, **kwargs):
+                """
+                Start the server.
+
+                :param kwargs: keyword arguments
+                """
+                await self.main_loop()
+
+        if 'CICD_JINA_DISABLE_HEALTHCHECK_LOGS' in os.environ:
+
+            class _EndpointFilter(logging.Filter):
+                def filter(self, record: logging.LogRecord) -> bool:
+                    # NOTE: space is important after `GET /`, else all logs will be disabled.
+                    return record.getMessage().find("GET / ") == -1
+
+            # Filter out healthcheck endpoint `GET /`
+            logging.getLogger("uvicorn.access").addFilter(_EndpointFilter())
+
+        uvicorn_kwargs = self.uvicorn_kwargs or {}
+
+        if self.ssl_keyfile and 'ssl_keyfile' not in uvicorn_kwargs.keys():
+            uvicorn_kwargs['ssl_keyfile'] = self.ssl_keyfile
+
+        if self.ssl_certfile and 'ssl_certfile' not in uvicorn_kwargs.keys():
+            uvicorn_kwargs['ssl_certfile'] = self.ssl_certfile
+
+        self.server = UviServer(
+            config=Config(
+                app=self.app,
+                host=self.host,
+                port=self.port,
+                log_level=os.getenv('JINA_LOG_LEVEL', 'error').lower(),
+                **uvicorn_kwargs,
+            )
+        )
+
+        await self.server.setup()
+
+    async def shutdown(self):
+        """
+        Free resources allocated when setting up HTTP server
+        """
+        self.server.should_exit = True
+        await self.server.shutdown()
+
+    async def run_server(self):
+        """Run HTTP server forever"""
+        await self.server.serve()
diff --git a/jina/serve/runtimes/gateway/http/gateway.py b/jina/serve/runtimes/gateway/http/gateway.py
index fd8d0311f859f..5b114f75766ec 100644
--- a/jina/serve/runtimes/gateway/http/gateway.py
+++ b/jina/serve/runtimes/gateway/http/gateway.py
@@ -2,13 +2,12 @@
 import os
 from typing import Optional
 
-from jina import __default_host__
 from jina.importer import ImportExtensions
-from jina.serve.gateway import BaseGateway
 from jina.serve.runtimes.gateway.http.app import get_fastapi_app
+from jina.serve.runtimes.gateway.http.fastapi import FastAPIBaseGateway
 
 
-class HTTPGateway(BaseGateway):
+class HTTPGateway(FastAPIBaseGateway):
     """HTTP Gateway implementation"""
 
     def __init__(
@@ -20,10 +19,6 @@ def __init__(
         expose_endpoints: Optional[str] = None,
         expose_graphql_endpoint: Optional[bool] = False,
         cors: Optional[bool] = False,
-        ssl_keyfile: Optional[str] = None,
-        ssl_certfile: Optional[str] = None,
-        uvicorn_kwargs: Optional[dict] = None,
-        proxy: Optional[bool] = None,
         **kwargs
     ):
         """Initialize the gateway
@@ -37,11 +32,6 @@ def __init__(
         :param expose_endpoints: A JSON string that represents a map from executor endpoints (`@requests(on=...)`) to HTTP endpoints.
         :param expose_graphql_endpoint: If set, /graphql endpoint is added to HTTP interface.
         :param cors: If set, a CORS middleware is added to FastAPI frontend to allow cross-origin access.
-        :param ssl_keyfile: the path to the key file
-        :param ssl_certfile: the path to the certificate file
-        :param uvicorn_kwargs: Dictionary of kwargs arguments that will be passed to Uvicorn server when starting the server
-        :param proxy: If set, respect the http_proxy and https_proxy environment variables, otherwise, it will unset
-            these proxy variables before start. gRPC seems to prefer no proxy
         :param kwargs: keyword args
         """
         super().__init__(**kwargs)
@@ -52,21 +42,17 @@ def __init__(
         self.expose_endpoints = expose_endpoints
         self.expose_graphql_endpoint = expose_graphql_endpoint
         self.cors = cors
-        self.ssl_keyfile = ssl_keyfile
-        self.ssl_certfile = ssl_certfile
-        self.uvicorn_kwargs = uvicorn_kwargs
 
-        if not proxy and os.name != 'nt':
-            os.unsetenv('http_proxy')
-            os.unsetenv('https_proxy')
-
-    async def setup_server(self):
+    @property
+    def app(self):
         """
-        Initialize and return GRPC server
+        FastAPI app needed to define an HTTP Jina Gateway
+
+        :return: FastAPI app
         """
         from jina.helper import extend_rest_interface
 
-        self.app = extend_rest_interface(
+        return extend_rest_interface(
             get_fastapi_app(
                 streamer=self.streamer,
                 title=self.title,
@@ -81,72 +67,3 @@ async def setup_server(self):
                 tracer_provider=self.tracer_provider,
             )
         )
-
-        with ImportExtensions(required=True):
-            from uvicorn import Config, Server
-
-        class UviServer(Server):
-            """The uvicorn server."""
-
-            async def setup(self, sockets=None):
-                """
-                Setup uvicorn server.
-
-                :param sockets: sockets of server.
-                """
-                config = self.config
-                if not config.loaded:
-                    config.load()
-                self.lifespan = config.lifespan_class(config)
-                await self.startup(sockets=sockets)
-                if self.should_exit:
-                    return
-
-            async def serve(self, **kwargs):
-                """
-                Start the server.
-
-                :param kwargs: keyword arguments
-                """
-                await self.main_loop()
-
-        if 'CICD_JINA_DISABLE_HEALTHCHECK_LOGS' in os.environ:
-
-            class _EndpointFilter(logging.Filter):
-                def filter(self, record: logging.LogRecord) -> bool:
-                    # NOTE: space is important after `GET /`, else all logs will be disabled.
-                    return record.getMessage().find("GET / ") == -1
-
-            # Filter out healthcheck endpoint `GET /`
-            logging.getLogger("uvicorn.access").addFilter(_EndpointFilter())
-
-        uvicorn_kwargs = self.uvicorn_kwargs or {}
-
-        if self.ssl_keyfile and 'ssl_keyfile' not in uvicorn_kwargs.keys():
-            uvicorn_kwargs['ssl_keyfile'] = self.ssl_keyfile
-
-        if self.ssl_certfile and 'ssl_certfile' not in uvicorn_kwargs.keys():
-            uvicorn_kwargs['ssl_certfile'] = self.ssl_certfile
-
-        self.server = UviServer(
-            config=Config(
-                app=self.app,
-                host=self.host,
-                port=self.port,
-                log_level=os.getenv('JINA_LOG_LEVEL', 'error').lower(),
-                **uvicorn_kwargs,
-            )
-        )
-
-        await self.server.setup()
-
-    async def shutdown(self):
-        """
-        Free resources allocated when setting up HTTP server
-        """
-        self.server.should_exit = True
-        await self.server.shutdown()
-
-    async def run_server(self):
-        """Run HTTP server forever"""
-        await self.server.serve()
diff --git a/tests/unit/orchestrate/flow/flow-construct/test_flow_custom_gateway.py b/tests/unit/orchestrate/flow/flow-construct/test_flow_custom_gateway.py
index 254ae855c626d..cccbf97144582 100644
--- a/tests/unit/orchestrate/flow/flow-construct/test_flow_custom_gateway.py
+++ b/tests/unit/orchestrate/flow/flow-construct/test_flow_custom_gateway.py
@@ -16,6 +16,10 @@
     cur_dir, '../../../yaml/test-custom-gateway.yml'
 )
 
+_dummy_fastapi_gateway_yaml_path = os.path.join(
+    cur_dir, '../../../yaml/test-fastapi-gateway.yml'
+)
+
 _flow_with_dummy_gateway_yaml_path = os.path.join(
     cur_dir, '../../../yaml/test-flow-custom-gateway-nested-config.yml'
 )
@@ -30,6 +34,11 @@
             {},
             {'arg1': 'hello', 'arg2': 'world', 'arg3': 'default-arg3'},
         ),
+        (
+            _dummy_fastapi_gateway_yaml_path,
+            {},
+            {'arg1': 'hello', 'arg2': 'world', 'arg3': 'default-arg3'},
+        ),
         (
             'DummyGateway',
             {'arg1': 'arg1', 'arg2': 'arg2', 'arg3': 'arg3'},
@@ -40,6 +49,11 @@
             {'arg1': 'arg1', 'arg2': 'arg2', 'arg3': 'arg3'},
             {'arg1': 'arg1', 'arg2': 'arg2', 'arg3': 'arg3'},
         ),
+        (
+            _dummy_fastapi_gateway_yaml_path,
+            {'arg1': 'arg1', 'arg2': 'arg2', 'arg3': 'arg3'},
+            {'arg1': 'arg1', 'arg2': 'arg2', 'arg3': 'arg3'},
+        ),
         (
             'DummyGateway',
             {'arg1': 'arg1'},
@@ -50,6 +64,11 @@
             {'arg1': 'arg1'},
             {'arg1': 'arg1', 'arg2': 'world', 'arg3': 'default-arg3'},
         ),
+        (
+            _dummy_fastapi_gateway_yaml_path,
+            {'arg1': 'arg1'},
+            {'arg1': 'arg1', 'arg2': 'world', 'arg3': 'default-arg3'},
+        ),
     ],
 )
 def test_flow_custom_gateway_no_executor(uses, uses_with, expected):
diff --git a/tests/unit/serve/gateway/test_gateway.py b/tests/unit/serve/gateway/test_gateway.py
index 07941fa3ad2fc..5722a8224b9d4 100644
--- a/tests/unit/serve/gateway/test_gateway.py
+++ b/tests/unit/serve/gateway/test_gateway.py
@@ -13,12 +13,13 @@
     _validate_custom_gateway_process,
     _validate_dummy_custom_gateway_response,
 )
-
 from tests.unit.yaml.dummy_gateway import DummyGateway
 
-
 cur_dir = os.path.dirname(os.path.abspath(__file__))
 _dummy_gateway_yaml_path = os.path.join(cur_dir, '../../yaml/test-custom-gateway.yml')
+_dummy_fastapi_gateway_yaml_path = os.path.join(
+    cur_dir, '../../yaml/test-fastapi-gateway.yml'
+)
 
 
 def _create_gateway_runtime(port, uses, uses_with, worker_port):
@@ -88,6 +89,11 @@ def _start_worker_runtime(uses):
             {},
             {'arg1': 'hello', 'arg2': 'world', 'arg3': 'default-arg3'},
         ),
+        (
+            _dummy_fastapi_gateway_yaml_path,
+            {},
+            {'arg1': 'hello', 'arg2': 'world', 'arg3': 'default-arg3'},
+        ),
         (
             'DummyGateway',
             {'arg1': 'arg1', 'arg2': 'arg2', 'arg3': 'arg3'},
@@ -98,6 +104,11 @@ def _start_worker_runtime(uses):
             {'arg1': 'arg1', 'arg2': 'arg2', 'arg3': 'arg3'},
             {'arg1': 'arg1', 'arg2': 'arg2', 'arg3': 'arg3'},
         ),
+        (
+            _dummy_fastapi_gateway_yaml_path,
+            {'arg1': 'arg1', 'arg2': 'arg2', 'arg3': 'arg3'},
+            {'arg1': 'arg1', 'arg2': 'arg2', 'arg3': 'arg3'},
+        ),
         (
             'DummyGateway',
             {'arg1': 'arg1'},
@@ -108,6 +119,11 @@ def _start_worker_runtime(uses):
             {'arg1': 'arg1'},
             {'arg1': 'arg1', 'arg2': 'world', 'arg3': 'default-arg3'},
         ),
+        (
+            _dummy_fastapi_gateway_yaml_path,
+            {'arg1': 'arg1'},
+            {'arg1': 'arg1', 'arg2': 'world', 'arg3': 'default-arg3'},
+        ),
     ],
 )
 def test_custom_gateway_no_executors(uses, uses_with, expected):
diff --git a/tests/unit/test_yamlparser.py b/tests/unit/test_yamlparser.py
index ae8559bf103eb..59e1e585ddc7f 100644
--- a/tests/unit/test_yamlparser.py
+++ b/tests/unit/test_yamlparser.py
@@ -170,23 +170,37 @@ def test_load_from_dict():
     assert b1.metas.name == 'hello123'
 
 
-def test_load_gateway_external_success():
+@pytest.mark.parametrize(
+    'yaml_file,gateway_name',
+    [
+        ('test-custom-gateway.yml', 'DummyGateway'),
+        ('test-fastapi-gateway.yml', 'DummyFastAPIGateway'),
+    ],
+)
+def test_load_gateway_external_success(yaml_file, gateway_name):
     with Gateway.load_config(
-        'yaml/test-custom-gateway.yml', runtime_args={'port': [12345]}
+        f'yaml/{yaml_file}', runtime_args={'port': [12345]}
     ) as gateway:
-        assert gateway.__class__.__name__ == 'DummyGateway'
+        assert gateway.__class__.__name__ == gateway_name
         assert gateway.arg1 == 'hello'
         assert gateway.arg2 == 'world'
         assert gateway.arg3 == 'default-arg3'
 
 
-def test_load_gateway_override_with():
+@pytest.mark.parametrize(
+    'yaml_file,gateway_name',
+    [
+        ('test-custom-gateway.yml', 'DummyGateway'),
+        ('test-fastapi-gateway.yml', 'DummyFastAPIGateway'),
+    ],
+)
+def test_load_gateway_override_with(yaml_file, gateway_name):
     with Gateway.load_config(
-        'yaml/test-custom-gateway.yml',
+        f'yaml/{yaml_file}',
         uses_with={'arg1': 'arg1', 'arg2': 'arg2', 'arg3': 'arg3'},
         runtime_args={'port': [12345]},
     ) as gateway:
-        assert gateway.__class__.__name__ == 'DummyGateway'
+        assert gateway.__class__.__name__ == gateway_name
         assert gateway.arg1 == 'arg1'
         assert gateway.arg2 == 'arg2'
         assert gateway.arg3 == 'arg3'
diff --git a/tests/unit/yaml/dummy_fastapi_gateway.py b/tests/unit/yaml/dummy_fastapi_gateway.py
new file mode 100644
index 0000000000000..1dac7d0c6d219
--- /dev/null
+++ b/tests/unit/yaml/dummy_fastapi_gateway.py
@@ -0,0 +1,61 @@
+from typing import Optional
+
+from docarray import Document, DocumentArray
+from pydantic import BaseModel
+
+from jina.clients.request import request_generator
+from jina.serve.runtimes.gateway.http.fastapi import FastAPIBaseGateway
+
+
+class DummyResponseModel(FastAPIBaseGateway):
+    arg1: Optional[str]
+    arg2: Optional[str]
+    arg3: Optional[str]
+
+
+class ProcessedResponseModel(BaseModel):
+    text: str
+    tags: Optional[dict]
+
+
+class DummyFastAPIGateway(FastAPIBaseGateway):
+    def __init__(
+        self, arg1: str = None, arg2: str = None, arg3: str = 'default-arg3', **kwargs
+    ):
+        super().__init__(**kwargs)
+        self.arg1 = arg1
+        self.arg2 = arg2
+        self.arg3 = arg3
+
+    @property
+    def app(self):
+        from fastapi import FastAPI
+
+        app = FastAPI(
+            title='Dummy Server',
+        )
+
+        @app.get(path='/', response_model=DummyResponseModel)
+        def _get_response():
+            return {
+                'arg1': self.arg1,
+                'arg2': self.arg2,
+                'arg3': self.arg3,
+            }
+
+        @app.get(
+            path='/stream',
+            response_model=ProcessedResponseModel,
+        )
+        async def _process(text: str):
+            doc = None
+            async for req in self.streamer.stream(
+                request_generator(
+                    exec_endpoint='/',
+                    data=DocumentArray([Document(text=text)]),
+                )
+            ):
+                doc = req.to_dict()['data'][0]
+            return {'text': doc['text'], 'tags': doc['tags']}
+
+        return app
diff --git a/tests/unit/yaml/test-fastapi-gateway.yml b/tests/unit/yaml/test-fastapi-gateway.yml
new file mode 100644
index 0000000000000..83d6b25500961
--- /dev/null
+++ b/tests/unit/yaml/test-fastapi-gateway.yml
@@ -0,0 +1,5 @@
+!DummyFastAPIGateway
+py_modules: dummy_fastapi_gateway.py
+with:
+  arg1: hello
+  arg2: world


Code review comments:
diff_hunk: @@ -0,0 +1,115 @@
+import logging
+import os
+from abc import abstractmethod
+from typing import Optional
+
+from jina.importer import ImportExtensions
+from jina.serve.gateway import BaseGateway
+
+
+class FastAPIBaseGateway(BaseGateway):
+    """HTTP Gateway implementation"""
+
+    def __init__(
+        self,
+        ssl_keyfile: Optional[str] = None,
+        ssl_certfile: Optional[str] = None,
+        uvicorn_kwargs: Optional[dict] = None,
+        proxy: Optional[bool] = None,
path: jina/serve/runtimes/gateway/http/fastapi.py
text: @user1:
a bool should not be Optional, should be `False` by default

diff_hunk: @@ -0,0 +1,115 @@
+import logging
+import os
+from abc import abstractmethod
+from typing import Optional
+
+from jina.importer import ImportExtensions
+from jina.serve.gateway import BaseGateway
+
+
+class FastAPIBaseGateway(BaseGateway):
+    """HTTP Gateway implementation"""
+
+    def __init__(
+        self,
+        ssl_keyfile: Optional[str] = None,
+        ssl_certfile: Optional[str] = None,
+        uvicorn_kwargs: Optional[dict] = None,
+        proxy: Optional[bool] = None,
+        **kwargs
+    ):
+        """Initialize the FastAPIBaseGateway
+        :param ssl_keyfile: the path to the key file
+        :param ssl_certfile: the path to the certificate file
+        :param uvicorn_kwargs: Dictionary of kwargs arguments that will be passed to Uvicorn server when starting the server
+        :param proxy: If set, respect the http_proxy and https_proxy environment variables, otherwise, it will unset
+            these proxy variables before start. gRPC seems to prefer no proxy
+        :param kwargs: keyword args
+        """
+        super().__init__(**kwargs)
+        self.ssl_keyfile = ssl_keyfile
path: jina/serve/runtimes/gateway/http/fastapi.py
text: @user1:
make them private