Repository: jina-ai/serve
Pull request: feat: add gateway API
As part of the Custom Gateway feature, we want to avoid confusion between gateway args and flow args.
We need to separate the concerns between gateway API and Flow API where the common parameter `uses` can be specified for both independently.
However, we do not want to introduce a breaking change and aim to keep existing gateway parameters at flow API and merge between those common gateway args and flow args.

The gateway API is added using the config_gateway method:
```python
f = Flow().config_gateway(uses='MyCustomGateway', port=12345)
```
and at yaml level:
```yaml
jtype: !Flow
gateway:
    uses: gateway.yml
    port: 12345
...
```
Since we don't want to introduce a breaking change, we aim to keep existing gateway parameters at flow API.
This means the gateway parameters inherit some valid Flow parameters and overrides them if specified (in the Gateway API):
```python
# ignores gateway params and uses gateway params instead
f = Flow(port=12344, protocol='http'').config_gateway(port=12345, protocol='grpc')
```

closes: https://github.com/jina-ai/jina/issues/5215


todos:

- [x] Flow.config_gateway
- [x] add gateway to json schema
- [x] update json schema tests
- [x] test custom gateway in Flow API (python class + yaml config + yaml and nested gateway yaml + docker ) in tests/unit/orchestrate/flow
- [ ] ~~fix the cli overload script and overload minimal args for each interface (flow) + add cli overload for the gateway method~~ document the inheritence of args from flow to executors behavior (tbd in https://github.com/jina-ai/jina/issues/5353)
- [x] Add `config_gateway` to cli overload script
- [x] yaml config tests in tests/unit/test_yaml_parser.py (yaml + nested gateway yaml)
- [x] Flow.save_config and Flow.load_config
- [x] Adapt k8s tests to new API
- [x] Adapt docker-compose tests to new API
- [x] check whether Flow API doesn't use Flow(uses=custom_gateway) anymore
- [x] pass CI
- [x] TODOs in PR

Problem statement:
Custom Gateway: expose gateway api in flow api as a separate set of args
As part of the Custom Gateway feature, we want to avoid confusion between gateway args and flow args.
We need to separate the concerns between gateway API and Flow API where the common parameter `uses` can be specified for both independently.
However, we do not want to introduce a breaking change and aim to keep existing gateway parameters at flow API and merge between those common gateway args and flow args in the following order:
1. gateway api args
2. flow api args
3 .default gateway api args
4. default flow api args
to be also added in flow yaml like this:
```jtype: Flow
... flow args
gateway:
    uses:
    ... # gateway args
executors:
    ...
```

Hints from issue discussion:

Commit message: test: test gateway in json schema

Patch to review:
diff --git a/jina/jaml/parsers/flow/v1.py b/jina/jaml/parsers/flow/v1.py
index 9c60d208dc5b7..3fd4531381607 100644
--- a/jina/jaml/parsers/flow/v1.py
+++ b/jina/jaml/parsers/flow/v1.py
@@ -79,7 +79,10 @@ def parse(
                 method = p_deployment_attr.get('method', 'add')
                 # support methods: add, needs, inspect
                 getattr(obj, method)(**p_deployment_attr, copy_flow=False)
-
+        gateway = data.get('gateway', {})
+        if gateway:
+            gateway_attr = {kk: expand_env_var(vv) for kk, vv in gateway.items()}
+            obj.config_gateway(**gateway_attr, copy_flow=False)
         return obj
 
     def dump(self, data: 'Flow') -> Dict:
@@ -111,8 +114,6 @@ def dump(self, data: 'Flow') -> Dict:
 
             # get nondefault kwargs
             parser = set_deployment_parser()
-            if v.role == DeploymentRoleType.GATEWAY:
-                parser = set_gateway_parser()
 
             non_default_kw = ArgNamespace.get_non_defaults_args(v.args, parser)
 
@@ -121,13 +122,21 @@ def dump(self, data: 'Flow') -> Dict:
             for t in _get_taboo(parser):
                 if t in kwargs:
                     kwargs.pop(t)
-            if k == 'gateway':
-                if 'JINA_FULL_CLI' in os.environ:
-                    r['with'].update(kwargs)
-                else:
-                    continue
-            else:
+            if k != 'gateway':
                 last_name = kwargs['name']
                 r['executors'].append(kwargs)
 
+        gateway_kwargs = {}
+        gateway_parser = set_gateway_parser()
+        non_default_kw = ArgNamespace.get_non_defaults_args(
+            data.gateway_args, gateway_parser
+        )
+        gateway_kwargs.update(non_default_kw)
+        for t in _get_taboo(gateway_parser):
+            if t in gateway_kwargs:
+                gateway_kwargs.pop(t)
+        if 'JINA_FULL_CLI' in os.environ:
+            r['with'].update(gateway_kwargs)
+        if gateway_kwargs:
+            r['gateway'] = gateway_kwargs
         return r
diff --git a/jina/orchestrate/flow/base.py b/jina/orchestrate/flow/base.py
index de4ad4f83b778..8a9787529a239 100644
--- a/jina/orchestrate/flow/base.py
+++ b/jina/orchestrate/flow/base.py
@@ -69,6 +69,7 @@
 from jina.orchestrate.deployments import Deployment
 from jina.orchestrate.flow.builder import _hanging_deployments, allowed_levels
 from jina.parsers import (
+    GATEWAY_NAME,
     set_client_cli_parser,
     set_deployment_parser,
     set_gateway_parser,
@@ -92,7 +93,6 @@ class FlowType(type(ExitStack), type(JAMLCompatible)):
     from jina.orchestrate.flow.asyncio import AsyncFlow
     from jina.serve.executors import BaseExecutor
 
-GATEWAY_NAME = 'gateway'
 FALLBACK_PARSERS = [
     set_gateway_parser(),
     set_deployment_parser(),
@@ -589,6 +589,14 @@ def _update_args(self, args, **kwargs):
         known_keys = vars(args)
         self._common_kwargs = {k: v for k, v in kwargs.items() if k not in known_keys}
 
+        # gateway args inherit from flow args
+        # TODO: we need a better exclusion of Flow parameters
+        self._gateway_kwargs = {
+            k: v
+            for k, v in self._common_kwargs.items()
+            if k not in {'uses', 'uses_with'}
+        }
+
         self._kwargs = ArgNamespace.get_non_defaults_args(
             args, _flow_parser
         )  #: for yaml dump
@@ -667,7 +675,7 @@ def _add_gateway(
         deployments_addresses: Dict[str, List[str]],
         deployments_metadata: Dict[str, Dict[str, str]],
         graph_conditions: Dict[str, Dict],
-        deployments_disabled_reduce: List[str],
+        deployments_disable_reduce: List[str],
         **kwargs,
     ):
         kwargs.update(
@@ -683,7 +691,7 @@ def _add_gateway(
             )
         )
 
-        kwargs.update(self._common_kwargs)
+        kwargs.update(self._gateway_kwargs)
         args = ArgNamespace.kwargs2namespace(kwargs, set_gateway_parser())
 
         # We need to check later if the port was manually set or randomly
@@ -695,7 +703,7 @@ def _add_gateway(
         args.graph_conditions = json.dumps(graph_conditions)
         args.deployments_addresses = json.dumps(deployments_addresses)
         args.deployments_metadata = json.dumps(deployments_metadata)
-        args.deployments_disable_reduce = json.dumps(deployments_disabled_reduce)
+        args.deployments_disable_reduce = json.dumps(deployments_disable_reduce)
         self._deployment_nodes[GATEWAY_NAME] = Deployment(args, needs)
 
     def _get_deployments_metadata(self) -> Dict[str, Dict[str, str]]:
@@ -1272,9 +1280,12 @@ def add(
         for key, value in op_flow._common_kwargs.items():
 
             # do not inherit from all the argument from the flow
+            # TODO: we should be explicit about args we inheret (maybe whitelist)
             if key not in kwargs and key not in [
                 'port',
                 'port_monitoring',
+                'uses',
+                'uses_with',
             ]:
                 kwargs[key] = value
 
@@ -1318,6 +1329,109 @@ def add(
 
         return op_flow
 
+    @allowed_levels([FlowBuildLevel.EMPTY])
+    def config_gateway(
+        self,
+        args: Optional['argparse.Namespace'] = None,
+        **kwargs,
+    ) -> Union['Flow', 'AsyncFlow']:
+
+        """
+        Configure the gateway
+        :param compression: The compression mechanism used when sending requests from the Head to the WorkerRuntimes. For more details, check https://grpc.github.io/grpc/python/grpc.html#compression.
+        :param connection_list: dictionary JSON with a list of connections to configure
+        :param disable_auto_volume: Do not automatically mount a volume for dockerized Executors.
+        :param disable_reduce: Disable the built-in reduce mechanism, set this if the reduction is to be handled by the Executor connected to this Head
+        :param docker_kwargs: Dictionary of kwargs arguments that will be passed to Docker SDK when starting the docker '
+          container.
+
+          More details can be found in the Docker SDK docs:  https://docker-py.readthedocs.io/en/stable/
+        :param entrypoint: The entrypoint command overrides the ENTRYPOINT in Docker image. when not set then the Docker image ENTRYPOINT takes effective.
+        :param env: The map of environment variables that are available inside runtime
+        :param exit_on_exceptions: List of exceptions that will cause the Gateway to shut down.
+        :param gpus: This argument allows dockerized Jina Gateway discover local gpu devices.
+
+              Note,
+              - To access all gpus, use `--gpus all`.
+              - To access multiple gpus, e.g. make use of 2 gpus, use `--gpus 2`.
+              - To access specified gpus based on device id, use `--gpus device=[YOUR-GPU-DEVICE-ID]`
+              - To access specified gpus based on multiple device id, use `--gpus device=[YOUR-GPU-DEVICE-ID1],device=[YOUR-GPU-DEVICE-ID2]`
+              - To specify more parameters, use `--gpus device=[YOUR-GPU-DEVICE-ID],runtime=nvidia,capabilities=display
+        :param grpc_metadata: The metadata to be passed to the gRPC request.
+        :param grpc_server_options: Dictionary of kwargs arguments that will be passed to the grpc server as options when starting the server, example : {'grpc.max_send_message_length': -1}
+        :param host_in: The host address for binding to, by default it is 0.0.0.0
+        :param install_requirements: If set, install `requirements.txt` in the Hub Gateway bundle to local
+        :param log_config: The YAML config of the logger used in this object.
+        :param name: The name of this object.
+
+              This will be used in the following places:
+              - how you refer to this object in Python/YAML/CLI
+              - visualization
+              - log message header
+              - ...
+
+              When not given, then the default naming strategy will apply.
+        :param port: The port for input data to bind to, default is a random port between [49152, 65535]. In the case of an external Executor (`--external` or `external=True`) this can be a list of ports, separated by commas. Then, every resulting address will be considered as one replica of the Executor.
+        :param port_monitoring: The port on which the prometheus server is exposed, default is a random port between [49152, 65535]
+        :param py_modules: The customized python modules need to be imported before loading the Gateway
+
+          Note that the recommended way is to only import a single module - a simple python file, if your
+          gateway can be defined in a single file, or an ``__init__.py`` file if you have multiple files,
+          which should be structured as a python package.
+        :param quiet: If set, then no log will be emitted from this object.
+        :param quiet_error: If set, then exception stack information will not be added to the log
+        :param quiet_remote_logs: Do not display the streaming of remote logs on local console
+        :param replicas: The number of replicas in the deployment
+        :param retries: Number of retries per gRPC call. If <0 it defaults to max(3, num_replicas)
+        :param shards: The number of shards in the deployment running at the same time. For more details check https://docs.jina.ai/fundamentals/flow/create-flow/#complex-flow-topologies
+        :param timeout_ctrl: The timeout in milliseconds of the control request, -1 for waiting forever
+        :param timeout_ready: The timeout in milliseconds of a Pod waits for the runtime to be ready, -1 for waiting forever
+        :param timeout_send: The timeout in milliseconds used when sending data requests to Executors, -1 means no timeout, disabled by default
+        :param tls: If set, connect to deployment using tls encryption
+        :param upload_files: The files on the host to be uploaded to the remote
+          workspace. This can be useful when your Deployment has more
+          file dependencies beyond a single YAML file, e.g.
+          Python files, data files.
+
+          Note,
+          - currently only flatten structure is supported, which means if you upload `[./foo/a.py, ./foo/b.pp, ./bar/c.yml]`, then they will be put under the _same_ workspace on the remote, losing all hierarchies.
+          - by default, `--uses` YAML file is always uploaded.
+          - uploaded files are by default isolated across the runs. To ensure files are submitted to the same workspace across different runs, use `--workspace-id` to specify the workspace.
+        :param uses: The config of the gateway, it could be one of the followings:
+                  * the string literal of an Gateway class name
+                  * a Gateway YAML file (.yml, .yaml, .jaml)
+                  * a docker image (must start with `docker://`)
+                  * the string literal of a YAML config (must start with `!` or `jtype: `)
+        :param uses_with: Dictionary of keyword arguments that will override the `with` configuration in `uses`
+        :param volumes: The path on the host to be mounted inside the container.
+
+          Note,
+          - If separated by `:`, then the first part will be considered as the local host path and the second part is the path in the container system.
+          - If no split provided, then the basename of that directory will be mounted into container's root path, e.g. `--volumes="/user/test/my-workspace"` will be mounted into `/my-workspace` inside the container.
+          - All volumes are mounted with read-write mode.
+        :param workspace: The working directory for any IO operations in this object. If not set, then derive from its parent `workspace`.
+        :param needs: the name of the Deployment(s) that this Deployment receives data from. One can also use "gateway" to indicate the connection with the gateway.
+        :param copy_flow: when set to true, then always copy the current Flow and do the modification on top of it then return, otherwise, do in-line modification
+        :param kwargs: other keyword-value arguments that the Deployment CLI supports
+        :return: a (new) Flow object with modification
+        :return: a (new) Flow object with modification
+
+        .. # noqa: DAR102
+        .. # noqa: DAR202
+        .. # noqa: DAR101
+        .. # noqa: DAR003
+        """
+
+        copy_flow = kwargs.pop('copy_flow', True)
+
+        op_flow = copy.deepcopy(self) if copy_flow else self
+
+        # TODO: properly merge between Flow and Gateway args
+        for key, value in kwargs.items():
+            op_flow._gateway_kwargs[key] = value
+
+        return op_flow
+
     @allowed_levels([FlowBuildLevel.EMPTY])
     def inspect(self, name: str = 'inspect', *args, **kwargs) -> 'Flow':
         """Add an inspection on the last changed Deployment in the Flow
@@ -1481,7 +1595,7 @@ def build(
                 deployments_addresses=op_flow._get_deployments_addresses(),
                 deployments_metadata=op_flow._get_deployments_metadata(),
                 graph_conditions=op_flow._get_graph_conditions(),
-                deployments_disabled_reduce=op_flow._get_disabled_reduce_deployments(),
+                deployments_disable_reduce=op_flow._get_disabled_reduce_deployments(),
                 uses=op_flow.args.uses,
             )
 
@@ -1776,7 +1890,7 @@ def client(self) -> 'BaseClient':
             port=self.port,
             protocol=self.protocol,
         )
-        kwargs.update(self._common_kwargs)
+        kwargs.update(self._gateway_kwargs)
         return Client(**kwargs)
 
     @property
@@ -1950,7 +2064,7 @@ def port(self) -> int:
         if GATEWAY_NAME in self._deployment_nodes:
             return self._deployment_nodes[GATEWAY_NAME].port
         else:
-            return self._common_kwargs.get('port', None)
+            return self._gateway_kwargs.get('port', None)
 
     @port.setter
     def port(self, value: int):
@@ -1958,11 +2072,11 @@ def port(self, value: int):
 
         :param value: the new port to expose
         """
-        self._common_kwargs['port'] = value
+        self._gateway_kwargs['port'] = value
 
         # Flow is build to graph already
         if self._build_level >= FlowBuildLevel.GRAPH:
-            self[GATEWAY_NAME].args.port = self._common_kwargs['port']
+            self[GATEWAY_NAME].args.port = self._gateway_kwargs['port']
 
         # Flow is running already, then close the existing gateway
         if self._build_level >= FlowBuildLevel.RUNNING:
@@ -1978,7 +2092,7 @@ def host(self) -> str:
         if GATEWAY_NAME in self._deployment_nodes:
             return self._deployment_nodes[GATEWAY_NAME].host
         else:
-            return self._common_kwargs.get('host', __default_host__)
+            return self._gateway_kwargs.get('host', __default_host__)
 
     @host.setter
     def host(self, value: str):
@@ -1986,11 +2100,11 @@ def host(self, value: str):
 
         :param value: the new port to expose
         """
-        self._common_kwargs['host'] = value
+        self._gateway_kwargs['host'] = value
 
         # Flow is build to graph already
         if self._build_level >= FlowBuildLevel.GRAPH:
-            self[GATEWAY_NAME].args.host = self._common_kwargs['host']
+            self[GATEWAY_NAME].args.host = self._gateway_kwargs['host']
 
         # Flow is running already, then close the existing gateway
         if self._build_level >= FlowBuildLevel.RUNNING:
@@ -2016,7 +2130,7 @@ def port_monitoring(self) -> Optional[int]:
         if GATEWAY_NAME in self._deployment_nodes:
             return self[GATEWAY_NAME].args.port_monitoring
         else:
-            return self._common_kwargs.get('port_monitoring', None)
+            return self._gateway_kwargs.get('port_monitoring', None)
 
     @property
     def address_private(self) -> str:
@@ -2217,7 +2331,7 @@ def protocol(self) -> GatewayProtocolType:
 
         :return: the protocol of this Flow
         """
-        v = self._common_kwargs.get('protocol', GatewayProtocolType.GRPC)
+        v = self._gateway_kwargs.get('protocol', GatewayProtocolType.GRPC)
         if isinstance(v, str):
             v = GatewayProtocolType.from_string(v)
         return v
@@ -2233,15 +2347,15 @@ def protocol(self, value: Union[str, GatewayProtocolType]):
             raise RuntimeError('Protocol can not be changed after the Flow has started')
 
         if isinstance(value, str):
-            self._common_kwargs['protocol'] = GatewayProtocolType.from_string(value)
+            self._gateway_kwargs['protocol'] = GatewayProtocolType.from_string(value)
         elif isinstance(value, GatewayProtocolType):
-            self._common_kwargs['protocol'] = value
+            self._gateway_kwargs['protocol'] = value
         else:
             raise TypeError(f'{value} must be either `str` or `GatewayProtocolType`')
 
         # Flow is build to graph already
         if self._build_level >= FlowBuildLevel.GRAPH:
-            self[GATEWAY_NAME].args.protocol = self._common_kwargs['protocol']
+            self[GATEWAY_NAME].args.protocol = self._gateway_kwargs['protocol']
 
     def __getitem__(self, item):
         if isinstance(item, str):
@@ -2520,9 +2634,9 @@ def client_args(self) -> argparse.Namespace:
 
         # noqa: DAR201
         """
-        if 'port' in self._common_kwargs:
-            kwargs = copy.deepcopy(self._common_kwargs)
-            kwargs['port'] = self._common_kwargs['port']
+        if 'port' in self._gateway_kwargs:
+            kwargs = copy.deepcopy(self._gateway_kwargs)
+            kwargs['port'] = self._gateway_kwargs['port']
 
         return ArgNamespace.kwargs2namespace(kwargs, set_client_cli_parser())
 
@@ -2532,14 +2646,14 @@ def gateway_args(self) -> argparse.Namespace:
 
         # noqa: DAR201
         """
-        return ArgNamespace.kwargs2namespace(self._common_kwargs, set_gateway_parser())
+        return ArgNamespace.kwargs2namespace(self._gateway_kwargs, set_gateway_parser())
 
     def _update_network_interface(self, **kwargs):
         """Update the network interface of this Flow (affects Gateway & Client)
 
         :param kwargs: new network settings
         """
-        self._common_kwargs.update(kwargs)
+        self._gateway_kwargs.update(kwargs)
 
     def __getattribute__(self, item):
         obj = super().__getattribute__(item)
diff --git a/jina/parsers/__init__.py b/jina/parsers/__init__.py
index 3b9990b8c3d4b..590db1526c95e 100644
--- a/jina/parsers/__init__.py
+++ b/jina/parsers/__init__.py
@@ -3,6 +3,8 @@
 from jina.parsers.orchestrate.runtimes.container import mixin_container_runtime_parser
 from jina.parsers.orchestrate.runtimes.head import mixin_head_parser
 
+GATEWAY_NAME = 'gateway'
+
 
 def set_pod_parser(parser=None):
     """Set the parser for the Pod
@@ -93,7 +95,7 @@ def set_gateway_parser(parser=None):
     from jina.enums import DeploymentRoleType
 
     parser.set_defaults(
-        name='gateway',
+        name=GATEWAY_NAME,
         runtime_cls='GatewayRuntime',
         deployment_role=DeploymentRoleType.GATEWAY,
     )
diff --git a/jina/schemas/__init__.py b/jina/schemas/__init__.py
index cdd8100c05b85..a1203c2dabd6a 100644
--- a/jina/schemas/__init__.py
+++ b/jina/schemas/__init__.py
@@ -1,3 +1,6 @@
+from jina.schemas.gateway import schema_gateway
+
+
 def get_full_schema() -> dict:
     """Get full schema
     :return: the full schema for Jina core as a dict.
@@ -16,6 +19,7 @@ def get_full_schema() -> dict:
         schema_metas,
         schema_deployment,
         IMPORTED.schema_executors,
+        schema_gateway,
     ]:
         definitions.update(s)
 
diff --git a/jina/schemas/flow.py b/jina/schemas/flow.py
index 65e8f555f42ea..3a6333220c48e 100644
--- a/jina/schemas/flow.py
+++ b/jina/schemas/flow.py
@@ -1,5 +1,6 @@
 from jina.schemas.helper import _cli_to_schema
 from jina_cli.export import api_to_dict
+from jina.schemas.gateway import schema_gateway
 
 _schema_flow_with = _cli_to_schema(
     api_to_dict(),
@@ -32,6 +33,7 @@
                 'items': {'$ref': '#/definitions/Jina::Deployment'},
                 'minItems': 1,
             },
+            'gateway': schema_gateway,
         },
         'type': 'object',
         'additionalProperties': False,
diff --git a/jina/schemas/gateway.py b/jina/schemas/gateway.py
new file mode 100644
index 0000000000000..8e3198fc579ec
--- /dev/null
+++ b/jina/schemas/gateway.py
@@ -0,0 +1,9 @@
+from jina.schemas.helper import _cli_to_schema
+from jina_cli.export import api_to_dict
+
+schema_gateway = _cli_to_schema(
+    api_to_dict(),
+    ['gateway'],
+    allow_addition=False,
+    description='The config of a Jina Gateway. A Gateway is a pod that encapsulates Flow logic and exposes services to the internet.',
+)['Jina::Gateway']
diff --git a/tests/docker_compose/test_docker_compose.py b/tests/docker_compose/test_docker_compose.py
index 5c09817104493..4e7791f207b42 100644
--- a/tests/docker_compose/test_docker_compose.py
+++ b/tests/docker_compose/test_docker_compose.py
@@ -336,14 +336,17 @@ async def test_flow_with_workspace(logger, docker_images, tmpdir):
     indirect=True,
 )
 async def test_flow_with_custom_gateway(logger, docker_images, tmpdir):
-    flow = Flow(
-        name='docker-compose-flow-custom-gateway',
-        port=9090,
-        protocol='http',
-        uses=f'docker://{docker_images[0]}',
-    ).add(
-        name='test_executor',
-        uses=f'docker://{docker_images[1]}',
+    flow = (
+        Flow(name='docker-compose-flow-custom-gateway')
+        .config_gateway(
+            port=9090,
+            protocol='http',
+            uses=f'docker://{docker_images[0]}',
+        )
+        .add(
+            name='test_executor',
+            uses=f'docker://{docker_images[1]}',
+        )
     )
 
     dump_path = os.path.join(str(tmpdir), 'docker-compose-flow-custom-gateway.yml')
diff --git a/tests/helper.py b/tests/helper.py
index 8dbd2815adea3..6ed60415fba2a 100644
--- a/tests/helper.py
+++ b/tests/helper.py
@@ -1,3 +1,16 @@
+from docarray import DocumentArray
+
+from jina import Executor, requests
+
+
+class ProcessExecutor(Executor):
+    @requests(on='/')
+    def process(self, docs: DocumentArray, **kwargs):
+        for doc in docs:
+            doc.text = doc.text + 'world'
+            doc.tags['processed'] = True
+
+
 def _validate_dummy_custom_gateway_response(port, expected):
     import requests
 
diff --git a/tests/k8s/test_k8s.py b/tests/k8s/test_k8s.py
index f1916240109d4..81ca254fa5a38 100644
--- a/tests/k8s/test_k8s.py
+++ b/tests/k8s/test_k8s.py
@@ -6,29 +6,27 @@
 import pytest
 import requests as req
 import yaml
+from docarray import DocumentArray
 from pytest_kind import cluster
 
-from docarray import DocumentArray
 from jina import Document, Executor, Flow, requests
 from jina.orchestrate.deployments import Deployment
 from jina.orchestrate.deployments.config.k8s import K8sDeploymentConfig
 from jina.parsers import set_deployment_parser
 from jina.serve.networking import GrpcConnectionPool
-from tests.helper import (
-    _validate_dummy_custom_gateway_response,
-)
+from tests.helper import _validate_dummy_custom_gateway_response
 
 cluster.KIND_VERSION = 'v0.11.1'
 
 
 async def create_all_flow_deployments_and_wait_ready(
-        flow_dump_path,
-        namespace,
-        api_client,
-        app_client,
-        core_client,
-        deployment_replicas_expected,
-        logger,
+    flow_dump_path,
+    namespace,
+    api_client,
+    app_client,
+    core_client,
+    deployment_replicas_expected,
+    logger,
 ):
     from kubernetes import utils
 
@@ -74,8 +72,8 @@ async def create_all_flow_deployments_and_wait_ready(
     while True:
         namespaced_pods = core_client.list_namespaced_pod(namespace)
         if (
-                namespaced_pods.items is not None
-                and len(namespaced_pods.items) == expected_deployments
+            namespaced_pods.items is not None
+            and len(namespaced_pods.items) == expected_deployments
         ):
             break
         logger.info(
@@ -103,8 +101,8 @@ async def create_all_flow_deployments_and_wait_ready(
                 )
             expected_num_replicas = deployment_replicas_expected[deployment_name]
             if (
-                    api_response.status.ready_replicas is not None
-                    and api_response.status.ready_replicas == expected_num_replicas
+                api_response.status.ready_replicas is not None
+                and api_response.status.ready_replicas == expected_num_replicas
             ):
                 logger.info(f'Deployment {deployment_name} is now ready')
                 deployments_ready.append(deployment_name)
@@ -127,14 +125,14 @@ async def run_test(flow, core_client, namespace, endpoint, n_docs=10, request_si
         core_client.list_namespaced_pod(
             namespace=namespace, label_selector='app=gateway'
         )
-            .items[0]
-            .metadata.name
+        .items[0]
+        .metadata.name
     )
     config_path = os.environ['KUBECONFIG']
     import portforward
 
     with portforward.forward(
-            namespace, gateway_pod_name, int(flow.port), int(flow.port), config_path
+        namespace, gateway_pod_name, int(flow.port), int(flow.port), config_path
     ):
         client_kwargs = dict(
             host='localhost',
@@ -147,10 +145,10 @@ async def run_test(flow, core_client, namespace, endpoint, n_docs=10, request_si
         client.show_progress = True
         responses = []
         async for resp in client.post(
-                endpoint,
-                inputs=[Document() for _ in range(n_docs)],
-                request_size=request_size,
-                return_responses=True,
+            endpoint,
+            inputs=[Document() for _ in range(n_docs)],
+            request_size=request_size,
+            return_responses=True,
         ):
             responses.append(resp)
 
@@ -221,21 +219,21 @@ def k8s_flow_with_needs(docker_images):
             port=9090,
             protocol='http',
         )
-            .add(
+        .add(
             name='segmenter',
             uses=f'docker://{docker_images[0]}',
         )
-            .add(
+        .add(
             name='textencoder',
             uses=f'docker://{docker_images[0]}',
             needs='segmenter',
         )
-            .add(
+        .add(
             name='imageencoder',
             uses=f'docker://{docker_images[0]}',
             needs='segmenter',
         )
-            .add(
+        .add(
             name='merger',
             uses=f'docker://{docker_images[1]}',
             needs=['imageencoder', 'textencoder'],
@@ -287,16 +285,16 @@ async def test_flow_with_monitoring(logger, tmpdir, docker_images, port_generato
         core_client.list_namespaced_pod(
             namespace=namespace, label_selector='app=gateway'
         )
-            .items[0]
-            .metadata.name
+        .items[0]
+        .metadata.name
     )
 
     executor_pod_name = (
         core_client.list_namespaced_pod(
             namespace=namespace, label_selector='app=segmenter'
         )
-            .items[0]
-            .metadata.name
+        .items[0]
+        .metadata.name
     )
 
     port_monitoring = GrpcConnectionPool.K8S_PORT_MONITORING
@@ -304,7 +302,7 @@ async def test_flow_with_monitoring(logger, tmpdir, docker_images, port_generato
 
     for pod_name in [gateway_pod_name, executor_pod_name]:
         with portforward.forward(
-                namespace, pod_name, port, port_monitoring, config_path
+            namespace, pod_name, port, port_monitoring, config_path
         ):
             resp = req.get(f'http://localhost:{port}/')
             assert resp.status_code == 200
@@ -876,14 +874,19 @@ async def test_flow_with_failing_executor(logger, docker_images, tmpdir):
     indirect=True,
 )
 async def test_flow_with_custom_gateway(logger, docker_images, tmpdir):
-    flow = Flow(
-        name='flow_with_custom_gateway',
-        port=9090,
-        protocol='http',
-        uses=f'docker://{docker_images[0]}',
-    ).add(
-        name='test_executor',
-        uses=f'docker://{docker_images[1]}',
+    flow = (
+        Flow(
+            name='flow_with_custom_gateway',
+        )
+        .config_gateway(
+            port=9090,
+            protocol='http',
+            uses=f'docker://{docker_images[0]}',
+        )
+        .add(
+            name='test_executor',
+            uses=f'docker://{docker_images[1]}',
+        )
     )
 
     dump_path = os.path.join(str(tmpdir), 'k8s-flow-custom-gateway.yml')
@@ -912,14 +915,14 @@ async def test_flow_with_custom_gateway(logger, docker_images, tmpdir):
         core_client.list_namespaced_pod(
             namespace=namespace, label_selector='app=gateway'
         )
-            .items[0]
-            .metadata.name
+        .items[0]
+        .metadata.name
     )
     config_path = os.environ['KUBECONFIG']
     import portforward
 
     with portforward.forward(
-            namespace, gateway_pod_name, flow.port, flow.port, config_path
+        namespace, gateway_pod_name, flow.port, flow.port, config_path
     ):
         _validate_dummy_custom_gateway_response(
             flow.port, {'arg1': 'hello', 'arg2': 'world', 'arg3': 'default-arg3'}
@@ -942,21 +945,16 @@ async def test_flow_with_custom_gateway(logger, docker_images, tmpdir):
     'docker_images', [['test-stateful-executor', 'jinaai/jina']], indirect=True
 )
 @pytest.mark.parametrize('workspace_path', ['workspace_path'])
-async def test_flow_with_stateful_executor(docker_images, tmpdir, logger, workspace_path):
+async def test_flow_with_stateful_executor(
+    docker_images, tmpdir, logger, workspace_path
+):
     dump_path = os.path.join(str(tmpdir), 'test-flow-with-volumes')
     namespace = f'test-flow-with-volumes'.lower()
-    flow = (
-        Flow(
-            name='test-flow-with-volumes',
-            port=9090,
-            protocol='http',
-        )
-            .add(
-            name='statefulexecutor',
-            uses=f'docker://{docker_images[0]}',
-            workspace=f'{str(tmpdir)}/workspace_path',
-            volumes=str(tmpdir)
-        )
+    flow = Flow(name='test-flow-with-volumes', port=9090, protocol='http',).add(
+        name='statefulexecutor',
+        uses=f'docker://{docker_images[0]}',
+        workspace=f'{str(tmpdir)}/workspace_path',
+        volumes=str(tmpdir),
     )
     flow.to_kubernetes_yaml(dump_path, k8s_namespace=namespace)
 
diff --git a/tests/unit/orchestrate/flow/flow-construct/test_flow.py b/tests/unit/orchestrate/flow/flow-construct/test_flow.py
index 716e5d4e1c6ad..a9e7f83baf434 100644
--- a/tests/unit/orchestrate/flow/flow-construct/test_flow.py
+++ b/tests/unit/orchestrate/flow/flow-construct/test_flow.py
@@ -8,7 +8,7 @@
 from docarray.document.generators import from_ndarray
 
 from jina import Document, DocumentArray, Executor, Flow, __windows__, requests
-from jina.enums import FlowBuildLevel
+from jina.enums import FlowBuildLevel, GatewayProtocolType
 from jina.excepts import RuntimeFailToStart
 from jina.helper import random_identity
 from jina.orchestrate.deployments import BaseDeployment
@@ -596,6 +596,30 @@ def test_gateway_only_flows_no_error(capsys, protocol):
     assert not captured.err
 
 
+@pytest.mark.slow
+def test_load_flow_with_custom_gateway(tmpdir):
+    # flow params are overridden by gateway params here
+    f = (
+        Flow(protocol='grpc', port=12344)
+        .config_gateway(uses='HTTPGateway', protocol='http', port=12345)
+        .add(name='executor')
+    )
+
+    with f:
+        _validate_flow(f)
+
+    f.save_config(os.path.join(str(tmpdir), 'tmp.yml'))
+    f = Flow.load_config(os.path.join(str(tmpdir), 'tmp.yml'))
+
+    assert f.port == 12345
+    assert f.protocol == GatewayProtocolType.HTTP
+
+    with Flow.load_config(os.path.join(str(tmpdir), 'tmp.yml')) as f:
+        assert f.port == 12345
+        assert f.protocol == GatewayProtocolType.HTTP
+        _validate_flow(f)
+
+
 def _validate_flow(f):
     graph_dict = f._get_graph_representation()
     addresses = f._get_deployments_addresses()
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
new file mode 100644
index 0000000000000..c165e5c025b0a
--- /dev/null
+++ b/tests/unit/orchestrate/flow/flow-construct/test_flow_custom_gateway.py
@@ -0,0 +1,61 @@
+import os
+
+import pytest
+
+from jina import Flow
+from tests.helper import (
+    ProcessExecutor,
+    _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'
+)
+
+
+@pytest.mark.parametrize(
+    'uses,uses_with,expected',
+    [
+        ('DummyGateway', {}, {'arg1': None, 'arg2': None, 'arg3': 'default-arg3'}),
+        (
+            _dummy_gateway_yaml_path,
+            {},
+            {'arg1': 'hello', 'arg2': 'world', 'arg3': 'default-arg3'},
+        ),
+        (
+            'DummyGateway',
+            {'arg1': 'arg1', 'arg2': 'arg2', 'arg3': 'arg3'},
+            {'arg1': 'arg1', 'arg2': 'arg2', 'arg3': 'arg3'},
+        ),
+        (
+            _dummy_gateway_yaml_path,
+            {'arg1': 'arg1', 'arg2': 'arg2', 'arg3': 'arg3'},
+            {'arg1': 'arg1', 'arg2': 'arg2', 'arg3': 'arg3'},
+        ),
+        (
+            'DummyGateway',
+            {'arg1': 'arg1'},
+            {'arg1': 'arg1', 'arg2': None, 'arg3': 'default-arg3'},
+        ),
+        (
+            _dummy_gateway_yaml_path,
+            {'arg1': 'arg1'},
+            {'arg1': 'arg1', 'arg2': 'world', 'arg3': 'default-arg3'},
+        ),
+    ],
+)
+def test_flow_custom_gateway_no_executor(uses, uses_with, expected):
+
+    flow = (
+        Flow()
+        .config_gateway(uses=uses, uses_with=uses_with)
+        .add(uses='ProcessExecutor')
+    )
+    with flow:
+        _validate_dummy_custom_gateway_response(flow.port, expected)
+        _validate_custom_gateway_process(
+            flow.port, 'hello', {'text': 'helloworld', 'tags': {'processed': True}}
+        )
diff --git a/tests/unit/orchestrate/flow/flow-construct/test_flow_yaml_parser.py b/tests/unit/orchestrate/flow/flow-construct/test_flow_yaml_parser.py
index 451f7716c3b2d..579930c151240 100644
--- a/tests/unit/orchestrate/flow/flow-construct/test_flow_yaml_parser.py
+++ b/tests/unit/orchestrate/flow/flow-construct/test_flow_yaml_parser.py
@@ -65,8 +65,11 @@ def test_load_dump_load(tmpdir):
     f2.save_config(str(Path(tmpdir) / 'a1.yml'))
 
 
-def test_load_modify_dump_load(tmpdir):
-    f: Flow = Flow.load_config('yaml/flow-gateway.yml')
+@pytest.mark.parametrize(
+    'yaml_file', ['yaml/flow-gateway.yml', 'yaml/flow-gateway-api.yml']
+)
+def test_load_modify_dump_load(tmpdir, yaml_file):
+    f: Flow = Flow.load_config(yaml_file)
     # assert vars inside `with`
     assert f._kwargs['name'] == 'abc'
     assert f.port == 12345
@@ -76,7 +79,7 @@ def test_load_modify_dump_load(tmpdir):
     assert f._deployment_nodes['custom2'].args.uses == 'CustomExecutor2'
     assert int(f._deployment_nodes['custom2'].args.port) == 23456
 
-    # change args inside `with`
+    # change args inside the gateway configuration
     f.port = 12346
     f.protocol = GatewayProtocolType.WEBSOCKET
     # change executor args
@@ -198,3 +201,18 @@ def test_flow_yaml_override_with_protocol():
     assert f2.protocol == GatewayProtocolType.HTTP
     f3 = Flow.load_config(path, uses_with={'protocol': 'websocket'})
     assert f3.protocol == GatewayProtocolType.WEBSOCKET
+
+
+@pytest.mark.parametrize(
+    'yaml_file', ['yaml/flow_with_gateway.yml', 'yaml/flow_with_gateway_nested.yml']
+)
+def test_load_flow_with_gateway(yaml_file):
+    path = os.path.join(cur_dir.parent.parent.parent, yaml_file)
+    flow = Flow.load_config(
+        path,
+        uses_with={'protocol': 'grpc', 'port': 12345},
+    )
+    with flow:
+        # protocol and port are overridden by the gateway configuration
+        assert flow.protocol == GatewayProtocolType.HTTP
+        assert flow.port == 12344
diff --git a/tests/unit/orchestrate/flow/flow-construct/yaml/flow-gateway-api.yml b/tests/unit/orchestrate/flow/flow-construct/yaml/flow-gateway-api.yml
new file mode 100644
index 0000000000000..c6e4b591210d5
--- /dev/null
+++ b/tests/unit/orchestrate/flow/flow-construct/yaml/flow-gateway-api.yml
@@ -0,0 +1,16 @@
+jtype: Flow
+version: 1
+gateway:
+  port: 12345
+  protocol: http
+with:
+  name: abc
+  port: 12344
+  protocol: grpc
+executors:
+  - name: custom1
+    uses: jinahub://CustomExecutor1
+  - name: custom2
+    uses: CustomExecutor2
+    port: 23456
+    port_out: 34567
diff --git a/tests/unit/schemas/test_schema.py b/tests/unit/schemas/test_schema.py
index eb99997688891..fb47a290d03cd 100644
--- a/tests/unit/schemas/test_schema.py
+++ b/tests/unit/schemas/test_schema.py
@@ -4,4 +4,8 @@
 
 
 def test_full_schema():
-    Draft7Validator.check_schema(get_full_schema())
+    schema = get_full_schema()
+    Draft7Validator.check_schema(schema)
+    # assert jina concepts exist in definitions
+    for concept in ['gateway', 'flow', 'metas', 'deployment']:
+        assert f'Jina::{concept.capitalize()}' in schema['definitions']
diff --git a/tests/unit/serve/gateway/test_gateway.py b/tests/unit/serve/gateway/test_gateway.py
index 581f87125686b..739ca9f3001bf 100644
--- a/tests/unit/serve/gateway/test_gateway.py
+++ b/tests/unit/serve/gateway/test_gateway.py
@@ -5,13 +5,12 @@
 
 import pytest
 
-from docarray import DocumentArray
-from jina import Executor, requests
 from jina.helper import random_port
 from jina.parsers import set_gateway_parser, set_pod_parser
 from jina.serve.runtimes.gateway import GatewayRuntime
 from jina.serve.runtimes.worker import WorkerRuntime
 from tests.helper import (
+    ProcessExecutor,
     _validate_custom_gateway_process,
     _validate_dummy_custom_gateway_response,
 )
@@ -21,14 +20,6 @@
 _dummy_gateway_yaml_path = os.path.join(cur_dir, '../../yaml/test-custom-gateway.yml')
 
 
-class ProcessExecutor(Executor):
-    @requests(on='/')
-    def process(self, docs: DocumentArray, **kwargs):
-        for doc in docs:
-            doc.text = doc.text + 'world'
-            doc.tags['processed'] = True
-
-
 def _create_gateway_runtime(port, uses, uses_with, worker_port):
     graph_description = '{"start-gateway": ["pod0"], "pod0": ["end-gateway"]}'
     pod_addresses = f'{{"pod0": ["0.0.0.0:{worker_port}"]}}'
diff --git a/tests/unit/yaml/flow_with_gateway.yml b/tests/unit/yaml/flow_with_gateway.yml
new file mode 100644
index 0000000000000..6d8218ab061b6
--- /dev/null
+++ b/tests/unit/yaml/flow_with_gateway.yml
@@ -0,0 +1,14 @@
+!Flow
+with:
+  protocol: grpc
+  port: 12345
+gateway:
+  py_modules: dummy_gateway.py
+  uses: DummyGateway
+  with:
+    arg1: hello
+    arg2: world
+  protocol: http
+  port: 12344
+executors:
+  - name: exec
\ No newline at end of file


Code review comments:
diff_hunk: @@ -3,6 +3,8 @@
 from jina.parsers.orchestrate.runtimes.container import mixin_container_runtime_parser
 from jina.parsers.orchestrate.runtimes.head import mixin_head_parser
 
+GATEWAY_NAME = 'gateway'
path: jina/parsers/__init__.py
text: @user1:
add this to some `helper` and use it everywhere in the code

@author:
done