Repository: jina-ai/serve
Pull request: feat(k8s): kubernetes completion, wait for ready and close

Problem statement:
Kubernetes -  Flow usage as seemless as native Flow
When creating a K8s Flow, I would like to use it as a native Flow:
```

flow = Flow(
    name='test-flow',
    port_expose=8080,
    infrastructure='K8S',
).add(...)

flow.start()

flow.post(
    on='/index',
    inputs=[Document(text=str(i), embedding=np.random(128)) for i in range(100)],
    show_progress=True,
)
```

For this we need several sub-tasks:

- [x] automatic port-forwarding when the gateway is deployed and ready
- [x] Use K8sPod as context manager, implement proper `wait_start_success` and `close` logic
- [x] Add readinessProbe through control Request

Kubernetes - Support graceful deletion 
**Describe the feature**

It is hard to delete the flow now. Users have to delete the namespace directly to delete all resources.

Hints from issue discussion:
I think the port forwarding is questionable? That is a debug mechanism you should use while developing and not for production so I would not do that automatically from code? The rest is fine, especially the readyness check we need
Hmm I see. Would you be fine with an argument `automatic_port_forwarding` which defaults to `False`?
That would be an option, but for me that does not really belong in our code base conceptually speaking (not even sure its possible).
Port forwarding is some hack you do on our computer (like a SSH tunnel) to reach a remote entity in K8s. I think that hack should be up to the user and not be built into our application.
but I think what we actually should do is: Enable the user to change the service type for the gateway to LoadBalancer
This would be needed for a production deployment without portforwarding
I prefer the consistent `with` context manager experience as we taught people in Flow docs, so this feature (which I guess is only about a proper deconstructor) needs to be implemented.

https://github.com/jina-ai/jina/pull/3357#discussion_r703318054
the destructor we need as well, agreed.

But I think we should at least also look into the readiness checks in kubernetes and if they are actually doing anything meaningful atm. Although these are two completely different things so lets split them  into two tickets?
Readiness checks maybe can be done by us by using `Control port` to check for `status`
> Readiness checks maybe can be done by us by using `Control port` to check for `status`

I am wondering if we could use Kubernetes built-in readiness probe?
> > Readiness checks maybe can be done by us by using `Control port` to check for `status`
> 
> I am wondering if we could use Kubernetes built-in readiness probe?

Actually, the readiness probe of Kubernetes should be changed by default to check the status endpoint
 the point here is if jina local to the user would have access to this information
To clarify the scope. No need to use http for readiness. Readinesa should be checked by hitting the control port of the Runtime. Liveness will be equivalent but if a problem or bottleneck is encounteres we will find a solution
/kind feature


@gaocegege thx for reporting this issue :) I completely agree that this is a needed feature. Han also mentioned it here in the past https://github.com/jina-ai/jina/issues/3334#issuecomment-915853961

I will add it to our roadmap. I am a bit busy the next days unfortunately, but maybe somebody else can pick it up. If nobody does it, I will when I am available again.

Fyi @florian-hoenicke 
Thanks for raising the issue. It will be handled in this ticket:
https://github.com/jina-ai/jina/issues/3334

Commit message: test: remove session

Patch to review:
diff --git a/.github/workflows/cd.yml b/.github/workflows/cd.yml
index b743c87cde173..137087a673803 100644
--- a/.github/workflows/cd.yml
+++ b/.github/workflows/cd.yml
@@ -263,8 +263,8 @@ jobs:
           JINAHUB_PASSWORD: ${{ secrets.JINAHUB_PASSWORD }}
 
   k8s-test:
-    runs-on: ubuntu-latest
     needs: update-schema
+    runs-on: ubuntu-latest
     steps:
       - uses: actions/checkout@v2
       - name: Set up Python 3.7
@@ -283,8 +283,21 @@ jobs:
           GITHUB_TOKEN: ${{ secrets.GITHUB_TOKEN }}
       - name: Test k8s
         run: |
-          ./tests/k8s/test_k8s.sh
+          pytest -v -s --suppress-no-test-exit-code --force-flaky --min-passes 1 --max-runs 5 --cov=jina --cov-report=xml ./tests/k8s/
         timeout-minutes: 30
+      - name: Check codecov file
+        id: check_files
+        uses: andstor/file-existence-action@v1
+        with:
+          files: "coverage.xml"
+      - name: Upload coverage from test to Codecov
+        uses: codecov/codecov-action@v1
+        if: steps.check_files.outputs.files_exists == 'true' && ${{ matrix.python-version }} == '3.7'
+        with:
+          file: coverage.xml
+          name: ${{ matrix.test-path }}-codecov
+          flags: ${{ steps.test.outputs.codecov_flag }}
+          fail_ci_if_error: false
 
   benchmark-pre-release:
     runs-on: ubuntu-latest
diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml
index 40105d690a595..e218ae1455eb2 100644
--- a/.github/workflows/ci.yml
+++ b/.github/workflows/ci.yml
@@ -223,8 +223,21 @@ jobs:
           GITHUB_TOKEN: ${{ secrets.GITHUB_TOKEN }}
       - name: Test k8s
         run: |
-          ./tests/k8s/test_k8s.sh
+          pytest -v -s --suppress-no-test-exit-code --force-flaky --min-passes 1 --max-runs 5 --cov=jina --cov-report=xml ./tests/k8s/test_k8s.py::test_flow_with_sharding
         timeout-minutes: 30
+      - name: Check codecov file
+        id: check_files
+        uses: andstor/file-existence-action@v1
+        with:
+          files: "coverage.xml"
+      - name: Upload coverage from test to Codecov
+        uses: codecov/codecov-action@v1
+        if: steps.check_files.outputs.files_exists == 'true' && ${{ matrix.python-version }} == '3.7'
+        with:
+          file: coverage.xml
+          name: ${{ matrix.test-path }}-codecov
+          flags: ${{ steps.test.outputs.codecov_flag }}
+          fail_ci_if_error: false
 
   prep-testbed:
     runs-on: ubuntu-latest
diff --git a/extra-requirements.txt b/extra-requirements.txt
index 74c19586ae2ad..051f311829886 100644
--- a/extra-requirements.txt
+++ b/extra-requirements.txt
@@ -73,6 +73,6 @@ dgl:                        cicd
 black==20.8b1:              test
 matplotlib:                 devel
 kubernetes>=18.20.0:        perf, standard, devel
-pykube-ng==21.3.0:          test
 pytest-kind==21.1.3:        test
-pytest-lazy-fixture:        test
\ No newline at end of file
+pytest-lazy-fixture:        test
+portforward>=0.2.4:         test
\ No newline at end of file
diff --git a/jina/flow/base.py b/jina/flow/base.py
index 0873b7769b609..8189cc9ed55d9 100644
--- a/jina/flow/base.py
+++ b/jina/flow/base.py
@@ -46,6 +46,7 @@
 from ..parsers import set_gateway_parser, set_pod_parser, set_client_cli_parser
 from ..parsers.flow import set_flow_parser
 from ..peapods import CompoundPod, Pod
+from ..peapods.pods.k8s import K8sPod
 from ..peapods.pods.factory import PodFactory
 from ..types.routing.table import RoutingTable
 from ..peapods.networking import is_remote_local_connection
@@ -78,6 +79,37 @@ class FlowType(type(ExitStack), type(JAMLCompatible)):
 class Flow(PostMixin, JAMLCompatible, ExitStack, metaclass=FlowType):
     """Flow is how Jina streamlines and distributes Executors. """
 
+    class _FlowK8sInfraResourcesManager:
+        def __init__(self, k8s_namespace: str, k8s_custom_resource_dir: Optional[str]):
+            self.k8s_namespace = k8s_namespace
+            self.k8s_custom_resource_dir = k8s_custom_resource_dir
+            self.namespace_created = False
+
+        def __enter__(self):
+            from ..peapods.pods.k8slib import kubernetes_tools
+
+            client = kubernetes_tools.K8sClients().core_v1
+            list_namespaces = [
+                item.metadata.name for item in client.list_namespace().items
+            ]
+            if self.k8s_namespace not in list_namespaces:
+                with JinaLogger(f'create_{self.k8s_namespace}') as logger:
+                    logger.info(f'🏝️\tCreate Namespace "{self.k8s_namespace}"')
+                    kubernetes_tools.create(
+                        'namespace',
+                        {'name': self.k8s_namespace},
+                        logger=logger,
+                        custom_resource_dir=self.k8s_custom_resource_dir,
+                    )
+                    self.namespace_created = True
+
+        def __exit__(self, exc_type, exc_val, exc_tb):
+            from ..peapods.pods.k8slib import kubernetes_tools
+
+            if self.namespace_created:
+                client = kubernetes_tools.K8sClients().core_v1
+                client.delete_namespace(name=self.k8s_namespace)
+
     # overload_inject_start_client_flow
     @overload
     def __init__(
@@ -323,6 +355,14 @@ def __init__(
         ]  #: default first pod is gateway, will add when build()
         self._update_args(args, **kwargs)
 
+        self.k8s_infrastructure_manager = None
+        if self.args.infrastructure == InfrastructureType.K8S:
+            self.k8s_infrastructure_manager = self._FlowK8sInfraResourcesManager(
+                k8s_namespace=self.args.name,
+                k8s_custom_resource_dir=getattr(
+                    self.args, 'k8s_custom_resource_dir', None
+                ),
+            )
         if isinstance(self.args, argparse.Namespace):
             self.logger = JinaLogger(
                 self.__class__.__name__, **vars(self.args), **self._common_kwargs
@@ -893,21 +933,40 @@ def _get_routing_table(self) -> RoutingTable:
                     )
 
         for end, pod in self._pod_nodes.items():
+
             if end == GATEWAY_NAME:
                 end = f'end-{GATEWAY_NAME}'
 
             if pod.head_args.hosts_in_connect is None:
                 pod.head_args.hosts_in_connect = []
 
+            if isinstance(pod, K8sPod):
+                from ..peapods.pods.k8slib import kubernetes_deployment
+
+                end = kubernetes_deployment.to_dns_name(end)
             if end not in graph.pods:
                 end = end + '_head'
+            if isinstance(pod, K8sPod):
+                from ..peapods.pods.k8slib import kubernetes_deployment
+
+                end = kubernetes_deployment.to_dns_name(end)
 
             for start in pod.needs:
+                start_pod = self._pod_nodes[start]
+
                 if start == GATEWAY_NAME:
                     start = f'start-{GATEWAY_NAME}'
 
+                if isinstance(start_pod, K8sPod):
+                    from ..peapods.pods.k8slib import kubernetes_deployment
+
+                    start = kubernetes_deployment.to_dns_name(start)
                 if start not in graph.pods:
                     start = start + '_tail'
+                if isinstance(start_pod, K8sPod):
+                    from ..peapods.pods.k8slib import kubernetes_deployment
+
+                    start = kubernetes_deployment.to_dns_name(start)
 
                 start_pod = graph._get_target_pod(start)
 
@@ -1040,8 +1099,11 @@ def __exit__(self, exc_type, exc_val, exc_tb):
         if self.args.env:
             for k in self.args.env.keys():
                 os.environ.pop(k, None)
+
+        # do not know why but removing these 2 lines make 2 tests fail
         if GATEWAY_NAME in self._pod_nodes:
             self._pod_nodes.pop(GATEWAY_NAME)
+
         self._build_level = FlowBuildLevel.EMPTY
         self.logger.debug('Flow is closed!')
         self.logger.close()
@@ -1062,6 +1124,9 @@ def start(self):
         if self._build_level.value < FlowBuildLevel.GRAPH.value:
             self.build(copy_flow=False)
 
+        if self.k8s_infrastructure_manager is not None:
+            self.enter_context(self.k8s_infrastructure_manager)
+
         # set env only before the Pod get started
         if self.args.env:
             for k, v in self.args.env.items():
diff --git a/jina/peapods/networking.py b/jina/peapods/networking.py
index 6b286e08ffc50..fb14b9eb94836 100644
--- a/jina/peapods/networking.py
+++ b/jina/peapods/networking.py
@@ -439,12 +439,12 @@ def create_connection_pool(args: 'Namespace') -> ConnectionPool:
     :return: A connection pool object
     """
     if args.k8s_namespace and args.k8s_connection_pool:
-        from jina.peapods.pods.k8slib.kubernetes_tools import K8SClients
+        from jina.peapods.pods.k8slib.kubernetes_tools import K8sClients
 
-        k8s_clients = K8SClients()
+        k8s_clients = K8sClients()
         return K8sGrpcConnectionPool(
             namespace=args.k8s_namespace,
-            client=k8s_clients.v1,
+            client=k8s_clients.core_v1,
         )
     else:
         return GrpcConnectionPool()
diff --git a/jina/peapods/pods/__init__.py b/jina/peapods/pods/__init__.py
index 4abe3aa0be059..3e18b2173f837 100644
--- a/jina/peapods/pods/__init__.py
+++ b/jina/peapods/pods/__init__.py
@@ -483,7 +483,7 @@ def _activate(self):
 
         self._activated = True
 
-    def start(self) -> 'BasePod':
+    def start(self) -> 'Pod':
         """
         Start to run all :class:`BasePea` in this BasePod.
 
diff --git a/jina/peapods/pods/k8s.py b/jina/peapods/pods/k8s.py
index 060c18faef20b..b51781809cdaf 100644
--- a/jina/peapods/pods/k8s.py
+++ b/jina/peapods/pods/k8s.py
@@ -1,18 +1,215 @@
 import copy
 import os
+import time
 from argparse import Namespace
 from typing import Optional, Dict, Union, Set, List
 
 import jina
 from .k8slib import kubernetes_deployment, kubernetes_tools
-from .. import BasePod
+from ..pods import BasePod, ExitFIFO
 from ... import __default_executor__
 from ...logging.logger import JinaLogger
+from ...excepts import RuntimeFailToStart
 
 
-class K8sPod(BasePod):
+class K8sPod(BasePod, ExitFIFO):
     """The K8sPod (KubernetesPod)  is used for deployments on Kubernetes."""
 
+    class _K8sDeployment:
+        def __init__(
+            self,
+            name: str,
+            head_port_in: int,
+            tail_port_out: int,
+            head_zmq_identity: bytes,
+            version: str,
+            shard_id: Optional[int],
+            common_args: Union['Namespace', Dict],
+            deployment_args: Union['Namespace', Dict],
+        ):
+            self.name = name
+            self.dns_name = kubernetes_deployment.to_dns_name(name)
+            self.head_port_in = head_port_in
+            self.tail_port_out = tail_port_out
+            self.head_zmq_identity = head_zmq_identity
+            self.version = version
+            self.shard_id = shard_id
+            self.common_args = common_args
+            self.deployment_args = deployment_args
+            self.k8s_namespace = self.common_args.k8s_namespace
+            self.num_replicas = getattr(self.deployment_args, 'replicas', 1)
+
+        def _deploy_gateway(self):
+            test_pip = os.getenv('JINA_K8S_USE_TEST_PIP') is not None
+            image_name = (
+                'jinaai/jina:test-pip'
+                if test_pip
+                else f'jinaai/jina:{self.version}-py38-standard'
+            )
+            kubernetes_deployment.deploy_service(
+                self.dns_name,
+                namespace=self.k8s_namespace,
+                image_name=image_name,
+                container_cmd='["jina"]',
+                container_args=f'["gateway", '
+                f'"--grpc-data-requests", '
+                f'{kubernetes_deployment.get_cli_params(self.common_args, ("pod_role",))}]',
+                logger=JinaLogger(f'deploy_{self.name}'),
+                replicas=1,
+                pull_policy='IfNotPresent',
+                port_expose=self.common_args.port_expose,
+            )
+
+        @staticmethod
+        def _construct_runtime_container_args(
+            deployment_args, uses, uses_metas, uses_with_string
+        ):
+            container_args = (
+                f'["executor", '
+                f'"--native", '
+                f'"--uses", "{uses}", '
+                f'"--grpc-data-requests", '
+                f'"--runtime-cls", "GRPCDataRuntime", '
+                f'"--uses-metas", "{uses_metas}", '
+                + uses_with_string
+                + f'{kubernetes_deployment.get_cli_params(deployment_args)}]'
+            )
+            return container_args
+
+        def _deploy_runtime(self):
+            image_name = kubernetes_deployment.get_image_name(self.deployment_args.uses)
+            init_container_args = kubernetes_deployment.get_init_container_args(
+                self.common_args
+            )
+            uses_metas = kubernetes_deployment.dictionary_to_cli_param(
+                {'pea_id': self.shard_id}
+            )
+            uses_with = kubernetes_deployment.dictionary_to_cli_param(
+                self.deployment_args.uses_with
+            )
+            uses_with_string = f'"--uses-with", "{uses_with}", ' if uses_with else ''
+            if image_name == __default_executor__:
+                test_pip = os.getenv('JINA_K8S_USE_TEST_PIP') is not None
+                image_name = (
+                    'jinaai/jina:test-pip'
+                    if test_pip
+                    else f'jinaai/jina:{self.version}-py38-standard'
+                )
+                uses = 'BaseExecutor'
+            else:
+                uses = 'config.yml'
+            container_args = self._construct_runtime_container_args(
+                self.deployment_args, uses, uses_metas, uses_with_string
+            )
+
+            kubernetes_deployment.deploy_service(
+                self.dns_name,
+                namespace=self.k8s_namespace,
+                image_name=image_name,
+                container_cmd='["jina"]',
+                container_args=container_args,
+                logger=JinaLogger(f'deploy_{self.name}'),
+                replicas=self.num_replicas,
+                pull_policy='IfNotPresent',
+                init_container=init_container_args,
+                custom_resource_dir=getattr(
+                    self.common_args, 'k8s_custom_resource_dir', None
+                ),
+            )
+
+        def wait_start_success(self):
+            _timeout = self.common_args.timeout_ready
+            if _timeout <= 0:
+                _timeout = None
+            else:
+                _timeout /= 1e3
+
+            from kubernetes import client
+
+            k8s_client = kubernetes_tools.K8sClients().apps_v1
+
+            with JinaLogger(f'waiting_for_{self.name}') as logger:
+                logger.info(
+                    f'🏝️\n\t\tWaiting for "{self.name}" to be ready, with {self.num_replicas} replicas'
+                )
+                timeout_ns = 1000000000 * _timeout if _timeout else None
+                now = time.time_ns()
+                exception_to_raise = None
+                while timeout_ns is None or time.time_ns() - now < timeout_ns:
+                    try:
+                        api_response = k8s_client.read_namespaced_deployment(
+                            name=self.dns_name, namespace=self.k8s_namespace
+                        )
+                        assert api_response.status.replicas == self.num_replicas
+                        if (
+                            api_response.status.ready_replicas is not None
+                            and api_response.status.ready_replicas == self.num_replicas
+                        ):
+                            logger.success(f' {self.name} has all its replicas ready!!')
+                            return
+                        else:
+                            ready_replicas = api_response.status.ready_replicas or 0
+                            logger.info(
+                                f'\nNumber of ready replicas {ready_replicas}, waiting for {self.num_replicas - ready_replicas} replicas to be available'
+                            )
+                            time.sleep(1.0)
+                    except client.ApiException as ex:
+                        exception_to_raise = ex
+                        break
+            fail_msg = f' Deployment {self.name} did not start with a timeout of {self.common_args.timeout_ready}'
+            if exception_to_raise:
+                fail_msg += f': {repr(exception_to_raise)}'
+            raise RuntimeFailToStart(fail_msg)
+
+        def start(self):
+            with JinaLogger(f'start_{self.name}') as logger:
+                logger.info(f'\t\tDeploying "{self.name}"')
+                if self.name == 'gateway':
+                    self._deploy_gateway()
+                else:
+                    self._deploy_runtime()
+                if not self.common_args.noblock_on_start:
+                    self.wait_start_success()
+            return self
+
+        def close(self):
+            from kubernetes import client
+
+            k8s_client = kubernetes_tools.K8sClients().apps_v1
+
+            with JinaLogger(f'close_{self.name}') as logger:
+                try:
+                    resp = k8s_client.delete_namespaced_deployment(
+                        name=self.dns_name, namespace=self.k8s_namespace
+                    )
+                    if resp.status == 'Success':
+                        logger.success(
+                            f' Successful deletion of deployment {self.name}'
+                        )
+                    else:
+                        logger.error(
+                            f' Deletion of deployment {self.name} unsuccessful with status {resp.status}'
+                        )
+                except client.ApiException as exc:
+                    logger.error(
+                        f' Error deleting deployment {self.name}: {exc.reason} '
+                    )
+
+        def __enter__(self):
+            return self.start()
+
+        def __exit__(self, exc_type, exc_val, exc_tb):
+            self.close()
+
+        def to_node(self):
+            return {
+                'name': self.dns_name,
+                'head_host': f'{self.dns_name}.{self.k8s_namespace}.svc',
+                'head_port_in': self.head_port_in,
+                'tail_port_out': self.tail_port_out,
+                'head_zmq_identity': self.head_zmq_identity,
+            }
+
     def __init__(
         self, args: Union['Namespace', Dict], needs: Optional[Set[str]] = None
     ):
@@ -24,6 +221,52 @@ def __init__(
 
         self.fixed_head_port_in = 8081
         self.fixed_tail_port_out = 8082
+        self.k8s_head_deployment = None
+        self.k8s_tail_deployment = None
+        if self.deployment_args['head_deployment'] is not None:
+            name = f'{self.name}-head'
+            self.k8s_head_deployment = self._K8sDeployment(
+                name=name,
+                head_port_in=self.fixed_head_port_in,
+                tail_port_out=self.fixed_tail_port_out,
+                head_zmq_identity=self.head_zmq_identity,
+                version=self.version,
+                shard_id=None,
+                common_args=self.args,
+                deployment_args=self.deployment_args['head_deployment'],
+            )
+        if self.deployment_args['tail_deployment'] is not None:
+            name = f'{self.name}-tail'
+            self.k8s_tail_deployment = self._K8sDeployment(
+                name=name,
+                head_port_in=self.fixed_head_port_in,
+                tail_port_out=self.fixed_tail_port_out,
+                head_zmq_identity=self.head_zmq_identity,
+                version=self.version,
+                shard_id=None,
+                common_args=self.args,
+                deployment_args=self.deployment_args['tail_deployment'],
+            )
+
+        self.k8s_deployments = []
+        for i, args in enumerate(self.deployment_args['deployments']):
+            name = (
+                f'{self.name}-{i}'
+                if len(self.deployment_args['deployments']) > 1
+                else f'{self.name}'
+            )
+            self.k8s_deployments.append(
+                self._K8sDeployment(
+                    name=name,
+                    head_port_in=self.fixed_head_port_in,
+                    tail_port_out=self.fixed_tail_port_out,
+                    head_zmq_identity=self.head_zmq_identity,
+                    version=self.version,
+                    shard_id=i,
+                    common_args=self.args,
+                    deployment_args=args,
+                )
+            )
 
     def _parse_args(
         self, args: Namespace
@@ -43,11 +286,13 @@ def _parse_deployment_args(self, args):
             # reasons to separate head and tail from peas is that they
             # can be deducted based on the previous and next pods
             parsed_args['head_deployment'] = copy.copy(args)
+            parsed_args['head_deployment'].replicas = 1
             parsed_args['head_deployment'].uses = (
                 args.uses_before or __default_executor__
             )
         if parallel > 1 or getattr(args, 'uses_after', None):
             parsed_args['tail_deployment'] = copy.copy(args)
+            parsed_args['tail_deployment'].replicas = 1
             parsed_args['tail_deployment'].uses = (
                 args.uses_after or __default_executor__
             )
@@ -56,7 +301,7 @@ def _parse_deployment_args(self, args):
         return parsed_args
 
     def __exit__(self, exc_type, exc_val, exc_tb) -> None:
-        self.close()
+        super().__exit__(exc_type, exc_val, exc_tb)
         self.join()
 
     @property
@@ -69,125 +314,40 @@ def host(self) -> str:
         """Not implemented"""
         raise NotImplementedError
 
-    def _deploy_gateway(self):
-        test_pip = os.getenv('JINA_K8S_USE_TEST_PIP') is not None
-        image_name = (
-            'jinaai/jina:test-pip'
-            if test_pip
-            else f'jinaai/jina:{self.version}-py38-standard'
-        )
-        kubernetes_deployment.deploy_service(
-            self.name,
-            namespace=self.args.k8s_namespace,
-            image_name=image_name,
-            container_cmd='["jina"]',
-            container_args=f'["gateway", '
-            f'"--grpc-data-requests", '
-            f'{kubernetes_deployment.get_cli_params(self.args, ("pod_role",))}]',
-            logger=JinaLogger(f'deploy_{self.name}'),
-            replicas=1,
-            pull_policy='IfNotPresent',
-            port_expose=self.args.port_expose,
-        )
-
-    def _deploy_runtime(self, deployment_args, replicas, deployment_id):
-        image_name = kubernetes_deployment.get_image_name(deployment_args.uses)
-        name_suffix = self.name + (
-            ''
-            if self.args.parallel == 1 and type(deployment_id) == int
-            else ('-' + str(deployment_id))
-        )
-        dns_name = kubernetes_deployment.to_dns_name(name_suffix)
-        init_container_args = kubernetes_deployment.get_init_container_args(self)
-        uses_metas = kubernetes_deployment.dictionary_to_cli_param(
-            {'pea_id': deployment_id}
-        )
-        uses_with = kubernetes_deployment.dictionary_to_cli_param(
-            deployment_args.uses_with
-        )
-        uses_with_string = f'"--uses-with", "{uses_with}", ' if uses_with else ''
-        if image_name == __default_executor__:
-            test_pip = os.getenv('JINA_K8S_USE_TEST_PIP') is not None
-            image_name = (
-                'jinaai/jina:test-pip'
-                if test_pip
-                else f'jinaai/jina:{self.version}-py38-standard'
-            )
-            uses = 'BaseExecutor'
-        else:
-            uses = 'config.yml'
-        container_args = self._construct_runtime_container_args(
-            deployment_args, uses, uses_metas, uses_with_string
-        )
-
-        kubernetes_deployment.deploy_service(
-            dns_name,
-            namespace=self.args.k8s_namespace,
-            image_name=image_name,
-            container_cmd='["jina"]',
-            container_args=container_args,
-            logger=JinaLogger(f'deploy_{self.name}'),
-            replicas=replicas,
-            pull_policy='IfNotPresent',
-            init_container=init_container_args,
-            custom_resource_dir=getattr(self.args, 'k8s_custom_resource_dir', None),
-        )
-
-    @staticmethod
-    def _construct_runtime_container_args(
-        deployment_args, uses, uses_metas, uses_with_string
-    ):
-        container_args = (
-            f'["executor", '
-            f'"--native", '
-            f'"--uses", "{uses}", '
-            f'"--grpc-data-requests", '
-            f'"--runtime-cls", "GRPCDataRuntime", '
-            f'"--uses-metas", "{uses_metas}", '
-            + uses_with_string
-            + f'{kubernetes_deployment.get_cli_params(deployment_args)}]'
-        )
-        return container_args
-
     def start(self) -> 'K8sPod':
         """Deploy the kubernetes pods via k8s Deployment and k8s Service.
 
         :return: self
         """
-        logger = JinaLogger(f'start_{self.name}')
-        logger.info(
-            f'🏝️\tCreate Namespace "{self.args.k8s_namespace}" for "{self.name}"'
-        )
-        kubernetes_tools.create(
-            'namespace',
-            {'name': self.args.k8s_namespace},
-            logger=logger,
-            custom_resource_dir=getattr(self.args, 'k8s_custom_resource_dir', None),
-        )
-        if self.name == 'gateway':
-            self._deploy_gateway()
-        else:
-            if self.deployment_args['head_deployment'] is not None:
-                self._deploy_runtime(self.deployment_args['head_deployment'], 1, 'head')
-
-            for i in range(self.args.parallel):
-                deployment_args = self.deployment_args['deployments'][i]
-                self._deploy_runtime(deployment_args, self.args.replicas, i)
-
-            if self.deployment_args['tail_deployment'] is not None:
-                self._deploy_runtime(self.deployment_args['tail_deployment'], 1, 'tail')
+        with JinaLogger(f'start_{self.name}') as logger:
+            logger.info(f'🏝️\tCreate deployments for "{self.name}"')
+            if self.k8s_head_deployment is not None:
+                self.enter_context(self.k8s_head_deployment)
+            for k8s_deployment in self.k8s_deployments:
+                self.enter_context(k8s_deployment)
+            if self.k8s_tail_deployment is not None:
+                self.enter_context(self.k8s_tail_deployment)
         return self
 
     def wait_start_success(self):
         """Not implemented. It should wait until the deployment is up and running"""
-        pass
-
-    def close(self):
-        """Not implemented. It should delete the namespace of the flow"""
-        pass
+        if not self.args.noblock_on_start:
+            raise ValueError(
+                f'{self.wait_start_success!r} should only be called when `noblock_on_start` is set to True'
+            )
+        try:
+            if self.k8s_head_deployment is not None:
+                self.k8s_head_deployment.wait_start_success()
+            for p in self.k8s_deployments:
+                p.wait_start_success()
+            if self.k8s_tail_deployment is not None:
+                self.k8s_tail_deployment.wait_start_success()
+        except:
+            self.close()
+            raise
 
     def join(self):
-        """Not implemented. It should wait to make sure deployments are properly killed."""
+        """Not needed. The context managers will manage the proper deletion"""
         pass
 
     def update_pea_args(self):
@@ -196,14 +356,6 @@ def update_pea_args(self):
         """
         self.deployment_args = self._parse_args(self.args)
 
-    @property
-    def is_ready(self) -> bool:
-        """Not implemented. It assumes it is ready.
-
-        :return: True
-        """
-        return True
-
     @property
     def head_args(self) -> Namespace:
         """Head args of the pod.
@@ -226,7 +378,19 @@ def num_peas(self) -> int:
 
         :return: number of peas
         """
-        return -1
+        return sum(
+            [
+                self.k8s_head_deployment.num_replicas
+                if self.k8s_head_deployment is not None
+                else 0
+            ]
+            + [
+                self.k8s_tail_deployment.num_replicas
+                if self.k8s_tail_deployment is not None
+                else 0
+            ]
+            + [k8s_deployment.num_replicas for k8s_deployment in self.k8s_deployments]
+        )
 
     @property
     def head_zmq_identity(self) -> bytes:
@@ -243,22 +407,15 @@ def deployments(self) -> List[Dict]:
         :return: list of dictionaries defining the attributes used by the routing table
         """
         res = []
+
         if self.args.name == 'gateway':
-            res.append(self._create_node(''))
+            res.append(self.k8s_deployments[0].to_node())
         else:
-            if self.deployment_args['head_deployment'] is not None:
-                res.append(self._create_node('head'))
-            for deployment_id, deployment_arg in enumerate(
-                self.deployment_args['deployments']
-            ):
-                name_suffix = (
-                    deployment_id
-                    if len(self.deployment_args['deployments']) > 1
-                    else ''
-                )
-                res.append(self._create_node(name_suffix))
-            if self.deployment_args['tail_deployment'] is not None:
-                res.append(self._create_node('tail'))
+            if self.k8s_head_deployment:
+                res.append(self.k8s_head_deployment.to_node())
+            res.extend([_.to_node() for _ in self.k8s_deployments])
+            if self.k8s_tail_deployment:
+                res.append(self.k8s_tail_deployment.to_node())
         return res
 
     def _get_base_executor_version(self):
@@ -272,17 +429,6 @@ def _get_base_executor_version(self):
         else:
             return 'master'
 
-    def _create_node(self, suffix):
-        name = f'{self.name}_{suffix}' if suffix != '' else self.name
-        dns_name = kubernetes_deployment.to_dns_name(name)
-        return {
-            'name': name,
-            'head_host': f'{dns_name}.{self.args.k8s_namespace}.svc',
-            'head_port_in': self.fixed_head_port_in,
-            'tail_port_out': self.fixed_tail_port_out,
-            'head_zmq_identity': self.head_zmq_identity,
-        }
-
     @property
     def _mermaid_str(self) -> List[str]:
         """String that will be used to represent the Pod graphically when `Flow.plot()` is invoked
diff --git a/jina/peapods/pods/k8slib/kubernetes_deployment.py b/jina/peapods/pods/k8slib/kubernetes_deployment.py
index 87e9afe9fd508..2f4e1f20f8cd5 100644
--- a/jina/peapods/pods/k8slib/kubernetes_deployment.py
+++ b/jina/peapods/pods/k8slib/kubernetes_deployment.py
@@ -199,18 +199,18 @@ def dictionary_to_cli_param(dictionary) -> str:
     return json.dumps(dictionary).replace('"', '\\"') if dictionary else ""
 
 
-def get_init_container_args(pod) -> Optional[Dict]:
+def get_init_container_args(args) -> Optional[Dict]:
     """Return the init container arguments for the k8s pod.
 
-    :param pod: pod where the init container is used.
+    :param args: args of the pod where the init container is used.
     :return: dictionary of init container arguments
     """
-    if pod.args.k8s_uses_init:
+    if args.k8s_uses_init:
         init_container = {
             'init-name': 'init',
-            'init-image': pod.args.k8s_uses_init,
-            'init-command': f'{pod.args.k8s_init_container_command}',
-            'mount-path': pod.args.k8s_mount_path,
+            'init-image': args.k8s_uses_init,
+            'init-command': f'{args.k8s_init_container_command}',
+            'mount-path': args.k8s_mount_path,
         }
     else:
         init_container = None
diff --git a/jina/peapods/pods/k8slib/kubernetes_tools.py b/jina/peapods/pods/k8slib/kubernetes_tools.py
index 3ea944f1956ce..e0cfc043e6eaf 100644
--- a/jina/peapods/pods/k8slib/kubernetes_tools.py
+++ b/jina/peapods/pods/k8slib/kubernetes_tools.py
@@ -1,30 +1,35 @@
 import os
 import tempfile
 import json
-from typing import Dict, Optional
+from typing import Dict, Optional, TYPE_CHECKING
 
-from jina.logging.logger import JinaLogger
-from jina.logging.predefined import default_logger
+from ....importer import ImportExtensions
+from ....logging.logger import JinaLogger
+from ....logging.predefined import default_logger
 
 cur_dir = os.path.dirname(__file__)
 DEFAULT_RESOURCE_DIR = os.path.join(
     cur_dir, '..', '..', '..', 'resources', 'k8s', 'template'
 )
 
+if TYPE_CHECKING:
+    from contextlib import GeneratorContextManager
 
-class K8SClients:
+
+class K8sClients:
     """
     The Kubernetes api is wrapped into a class to have a lazy reading of the cluster configuration.
 
     """
 
     def __init__(self):
-        self.__k8s_client = None
-        self.__v1 = None
-        self.__beta = None
-        self.__networking_v1_beta1_api = None
+        self._k8s_client = None
+        self._core_v1 = None
+        self._apps_v1 = None
+        self._beta = None
+        self._networking_v1_beta1_api = None
 
-    def __instantiate(self):
+    def _instantiate(self):
         # this import reads the `KUBECONFIG` env var. Lazy load to postpone the reading
         from kubernetes import config, client
 
@@ -36,11 +41,12 @@ def __instantiate(self):
             # this works if we are running inside k8s
             config.load_incluster_config()
 
-        self.__k8s_client = client.ApiClient()
-        self.__v1 = client.CoreV1Api(api_client=self.__k8s_client)
-        self.__beta = client.ExtensionsV1beta1Api(api_client=self.__k8s_client)
-        self.__networking_v1_beta1_api = client.NetworkingV1beta1Api(
-            api_client=self.__k8s_client
+        self._k8s_client = client.ApiClient()
+        self._core_v1 = client.CoreV1Api(api_client=self._k8s_client)
+        self._apps_v1 = client.AppsV1Api(api_client=self._k8s_client)
+        self._beta = client.ExtensionsV1beta1Api(api_client=self._k8s_client)
+        self._networking_v1_beta1_api = client.NetworkingV1beta1Api(
+            api_client=self._k8s_client
         )
 
     @property
@@ -49,19 +55,29 @@ def k8s_client(self):
 
         :return: k8s client
         """
-        if self.__k8s_client is None:
-            self.__instantiate()
-        return self.__k8s_client
+        if self._k8s_client is None:
+            self._instantiate()
+        return self._k8s_client
+
+    @property
+    def core_v1(self):
+        """V1 client for core
+
+        :return: v1 client
+        """
+        if self._core_v1 is None:
+            self._instantiate()
+        return self._core_v1
 
     @property
-    def v1(self):
+    def apps_v1(self):
         """V1 client for core
 
         :return: v1 client
         """
-        if self.__v1 is None:
-            self.__instantiate()
-        return self.__v1
+        if self._apps_v1 is None:
+            self._instantiate()
+        return self._apps_v1
 
     @property
     def beta(self):
@@ -69,9 +85,9 @@ def beta(self):
 
         :return: beta client
         """
-        if self.__beta is None:
-            self.__instantiate()
-        return self.__beta
+        if self._beta is None:
+            self._instantiate()
+        return self._beta
 
     @property
     def networking_v1_beta1_api(self):
@@ -79,12 +95,12 @@ def networking_v1_beta1_api(self):
 
         :return: networking client
         """
-        if self.__networking_v1_beta1_api is None:
-            self.__instantiate()
-        return self.__networking_v1_beta1_api
+        if self._networking_v1_beta1_api is None:
+            self._instantiate()
+        return self._networking_v1_beta1_api
 
 
-__k8s_clients = K8SClients()
+_k8s_clients = K8sClients()
 
 
 def create(
@@ -111,7 +127,7 @@ def create(
         with os.fdopen(fd, 'w') as tmp:
             tmp.write(yaml)
         try:
-            utils.create_from_yaml(__k8s_clients.k8s_client, path)
+            utils.create_from_yaml(_k8s_clients.k8s_client, path)
         except FailToCreateError as e:
             for api_exception in e.api_exceptions:
                 if api_exception.status == 409:
@@ -138,3 +154,36 @@ def _get_yaml(template: str, params: Dict, custom_resource_dir: Optional[str] =
         for k, v in params.items():
             content = content.replace(f'{{{k}}}', str(v))
     return content
+
+
+def _get_gateway_pod_name(namespace):
+    gateway_pod = _k8s_clients.core_v1.list_namespaced_pod(
+        namespace=namespace, label_selector='app=gateway'
+    )
+    return gateway_pod.items[0].metadata.name
+
+
+def get_port_forward_contextmanager(
+    namespace: str,
+    port_expose: int,
+    config_path: str = None,
+) -> 'GeneratorContextManager':
+    """Forward local requests to the gateway which is running in the Kubernetes cluster.
+    :param namespace: namespace of the gateway
+    :param port_expose: exposed port of the gateway
+    :param config_path: path to the Kubernetes config file
+    :return: context manager which sets up and terminates the port-forward
+    """
+    with ImportExtensions(
+        required=True,
+        help_text='sending requests to the Kubernetes cluster requires to install the portforward package, '
+        'please do `pip install jina[portforward]`',
+    ):
+        import portforward
+
+    gateway_pod_name = _get_gateway_pod_name(namespace)
+    if config_path is None and 'KUBECONFIG' in os.environ:
+        config_path = os.environ['KUBECONFIG']
+    return portforward.forward(
+        namespace, gateway_pod_name, port_expose, port_expose, config_path
+    )
diff --git a/jina/peapods/runtimes/grpc/__init__.py b/jina/peapods/runtimes/grpc/__init__.py
index e4d68cb95fe48..6ffe281c303d6 100644
--- a/jina/peapods/runtimes/grpc/__init__.py
+++ b/jina/peapods/runtimes/grpc/__init__.py
@@ -40,6 +40,7 @@ def __init__(self, args: argparse.Namespace, **kwargs):
         self._pending_tasks = []
         self._static_routing_table = args.static_routing_table
 
+        # Keep this initialization order, otherwise readiness check is not valid
         self._data_request_handler = DataRequestHandler(args, self.logger)
         self._grpclet = Grpclet(
             args=self.args,
diff --git a/jina/resources/extra-requirements.txt b/jina/resources/extra-requirements.txt
index 74c19586ae2ad..051f311829886 100644
--- a/jina/resources/extra-requirements.txt
+++ b/jina/resources/extra-requirements.txt
@@ -73,6 +73,6 @@ dgl:                        cicd
 black==20.8b1:              test
 matplotlib:                 devel
 kubernetes>=18.20.0:        perf, standard, devel
-pykube-ng==21.3.0:          test
 pytest-kind==21.1.3:        test
-pytest-lazy-fixture:        test
\ No newline at end of file
+pytest-lazy-fixture:        test
+portforward>=0.2.4:         test
\ No newline at end of file
diff --git a/jina/resources/k8s/template/deployment.yml b/jina/resources/k8s/template/deployment.yml
index 94b8f3919c98e..b446e5672ab14 100644
--- a/jina/resources/k8s/template/deployment.yml
+++ b/jina/resources/k8s/template/deployment.yml
@@ -25,11 +25,28 @@ spec:
           - containerPort: {port_in}
           - containerPort: {port_out}
           - containerPort: {port_ctrl}
+        readinessProbe:
+          exec:
+            command:
+              - python
+              - -c
+              - 'import os; import sys; from jina.peapods.runtimes.grpc import GRPCDataRuntime; control_port = int(os.getenv("JINA_EXECUTOR_CONTROL_PORT")); is_ready = GRPCDataRuntime.is_ready(ctrl_address=f"0.0.0.0:{str(control_port)}"); sys.exit(0) if is_ready else sys.exit(1)'
+          initialDelaySeconds: 5
+          periodSeconds: 5
+        livenessProbe:
+          exec:
+            command:
+              - python
+              - -c
+              - 'import os; import sys; from jina.peapods.runtimes.grpc import GRPCDataRuntime; control_port = int(os.getenv("JINA_EXECUTOR_CONTROL_PORT")); is_ready = GRPCDataRuntime.is_ready(ctrl_address=f"0.0.0.0:{str(control_port)}"); sys.exit(0) if is_ready else sys.exit(1)'
+          initialDelaySeconds: 30
+          periodSeconds: 60
         env:
-        - name: JINA_LOG_LEVEL
-          value: debug
-        - name: PYTHONUNBUFFERED
-          value: "1"
-        - name: WORKER_CLASS
-          value: "uvicorn.workers.UvicornH11Worker"
-
+          - name: JINA_LOG_LEVEL
+            value: DEBUG
+          - name: PYTHONUNBUFFERED
+            value: "1"
+          - name: WORKER_CLASS
+            value: "uvicorn.workers.UvicornH11Worker"
+          - name: JINA_EXECUTOR_CONTROL_PORT
+            value: "8081" # TODO there should be a better way to pass this and align with port_ctrl
diff --git a/tests/conftest.py b/tests/conftest.py
index 0483e475c3858..38ce670e823b8 100644
--- a/tests/conftest.py
+++ b/tests/conftest.py
@@ -191,7 +191,7 @@ def tmpfile(tmpdir):
     return tmpdir / tmpfile
 
 
-@pytest.fixture(scope="session")
+@pytest.fixture(scope='session')
 def event_loop(request):
     """
     Valid only for `pytest.mark.asyncio` tests
diff --git a/tests/k8s/conftest.py b/tests/k8s/conftest.py
index 66b2091227be2..4a5160affd588 100644
--- a/tests/k8s/conftest.py
+++ b/tests/k8s/conftest.py
@@ -1,11 +1,9 @@
-import operator
 import os
-from typing import List
 
 import docker
-import pykube
 import pytest
 from pytest_kind import KindCluster
+from typing import Optional
 
 from jina.logging.logger import JinaLogger
 
@@ -13,6 +11,24 @@
 cur_dir = os.path.dirname(__file__)
 
 
+class KindClusterWrapper:
+    def __init__(self, kind_cluster: KindCluster, logger: JinaLogger):
+        self._cluster = kind_cluster
+        self._cluster.ensure_kubectl()
+        self._kube_config_path = os.path.join(
+            os.getcwd(), '.pytest-kind/pytest-kind/kubeconfig'
+        )
+        self._log = logger
+        self._set_kube_config()
+
+    def _set_kube_config(self):
+        self._log.debug(f'Setting KUBECONFIG to {self._kube_config_path}')
+        os.environ['KUBECONFIG'] = self._kube_config_path
+
+    def load_docker_image(self, image_name: str):
+        self._cluster.load_docker_image(image_name)
+
+
 @pytest.fixture()
 def logger():
     logger = JinaLogger('kubernetes-testing')
@@ -24,6 +40,11 @@ def test_dir() -> str:
     return cur_dir
 
 
+@pytest.fixture()
+def k8s_cluster(kind_cluster: KindCluster, logger: JinaLogger) -> KindClusterWrapper:
+    return KindClusterWrapper(kind_cluster, logger)
+
+
 @pytest.fixture()
 def test_executor_image(logger: JinaLogger):
     image, build_logs = client.images.build(
@@ -60,79 +81,25 @@ def dummy_dumper_image(logger: JinaLogger):
     return image.tags[-1]
 
 
-class KindClusterWrapper:
-    def __init__(self, kind_cluster: KindCluster, logger: JinaLogger):
-        self._cluster = kind_cluster
-        self._cluster.ensure_kubectl()
-        self._kube_config_path = os.path.join(
-            os.getcwd(), '.pytest-kind/pytest-kind/kubeconfig'
-        )
-        self._kind_exec_path = os.path.join(
-            os.getcwd(), '.pytest-kind/pytest-kind/kind'
-        )
-        self._log = logger
-        self._set_kube_config()
-        self._pykube_api = self._cluster.api
-
-    def _set_kube_config(self):
-        self._log.debug(f'Seting KUBECONFIG to {self._kube_config_path}')
-        os.environ['KUBECONFIG'] = self._kube_config_path
-
-    @property
-    def kube_config_path(self) -> str:
-        return self._kube_config_path
-
-    @property
-    def kind_exec_path(self) -> str:
-        return self._kind_exec_path
-
-    def port_forward(
-        self,
-        service_name: str,
-        local_port: int,
-        service_port: int,
-        namespace: str = None,
-    ):
-        if namespace:
-            return self._cluster.port_forward(
-                service_name,
-                service_port,
-                '-n',
-                namespace,
-                local_port=local_port,
-                retries=20,
-            )
-        else:
-            return self._cluster.port_forward(
-                service_name, service_port, local_port=local_port, retries=20
-            )
-
-    def list_pods(self, namespace: str = None) -> List:
-        if namespace:
-            pod_list = list(pykube.Pod.objects(self._pykube_api, namespace=namespace))
-        else:
-            pod_list = list(pykube.Pod.objects(self._pykube_api))
-        return pod_list
-
-    def list_ready_pods(self, namespace: str = None) -> List:
-        return list(filter(operator.attrgetter("ready"), self.list_pods(namespace)))
-
-    def get_node_info(self):
-        nodes = []
-        for node in pykube.Node.objects(self._pykube_api):
-            nodes.append(node.obj)
-        return nodes
-
-    def needs_docker_image(self, image_name: str):
-        self._cluster.load_docker_image(image_name)
-
-
 @pytest.fixture()
-def k8s_cluster(kind_cluster: KindCluster, logger: JinaLogger) -> KindClusterWrapper:
-    yield KindClusterWrapper(kind_cluster, logger)
+def load_images_in_kind(
+    logger, test_executor_image, executor_merger_image, dummy_dumper_image, k8s_cluster
+):
+    logger.debug(f'Loading docker image into kind cluster...')
+    for image in [
+        test_executor_image,
+        executor_merger_image,
+        dummy_dumper_image,
+        'jinaai/jina:test-pip',
+    ]:
+        k8s_cluster.load_docker_image(image)
+    logger.debug(f'Done loading docker image into kind cluster...')
 
 
 @pytest.fixture()
-def k8s_cluster_namespaced(k8s_cluster) -> KindClusterWrapper:
-    yield k8s_cluster
-    k8s_cluster._cluster.kubectl('delete', 'namespace', 'test-flow')
+def set_test_pip_version():
+    import os
+
+    os.environ['JINA_K8S_USE_TEST_PIP'] = 'True'
+    yield
+    del os.environ['JINA_K8S_USE_TEST_PIP']
diff --git a/tests/k8s/test_custom_resource_dir.py b/tests/k8s/test_custom_resource_dir.py
index f817c37e50411..63e842246f1dd 100644
--- a/tests/k8s/test_custom_resource_dir.py
+++ b/tests/k8s/test_custom_resource_dir.py
@@ -1,45 +1,30 @@
-from unittest.mock import Mock
 import os
 
-from jina.peapods.pods.k8slib.kubernetes_deployment import kubernetes_tools
+from jina.peapods.pods.k8slib.kubernetes_tools import _get_yaml
 from jina import Flow
 
 
-def test_custom_resource_dir(monkeypatch):
-    mock_create = Mock()
-    monkeypatch.setattr(kubernetes_tools, 'create', mock_create)
+def test_custom_resource_dir():
     custom_resource_dir = '/test'
 
     flow = Flow(
         name='test-flow', port_expose=8080, infrastructure='K8S', protocol='http'
     ).add(name='test_executor', k8s_custom_resource_dir=custom_resource_dir)
-    flow.start()
-
-    # first three calls are for the executor the rest are calls for the gateway
-    for call in mock_create.call_args_list[:3]:
-        assert call[1]['custom_resource_dir'] == custom_resource_dir
-
+    assert (
+        flow._pod_nodes['test_executor'].args.k8s_custom_resource_dir
+        == custom_resource_dir
+    )
 
-def test_no_resource_dir_specified(monkeypatch):
-    mock_create = Mock()
-    monkeypatch.setattr(kubernetes_tools, 'create', mock_create)
 
+def test_no_resource_dir_specified():
     flow = Flow(
         name='test-flow', port_expose=8080, infrastructure='K8S', protocol='http'
     ).add(name='test_executor')
-    flow.start()
-
-    for call in mock_create.call_args_list:
-        assert (
-            'custom_resource_dir' not in call[1]
-            or call[1]['custom_resource_dir'] is None
-        )
+    assert flow._pod_nodes['test_executor'].args.k8s_custom_resource_dir is None
 
 
 def test_template_file_read_correctly(test_dir: str):
     custom_resource_dir = os.path.join(test_dir, 'custom-resource')
-    content = kubernetes_tools._get_yaml(
-        'namespace', params={}, custom_resource_dir=custom_resource_dir
-    )
+    content = _get_yaml('namespace', params={}, custom_resource_dir=custom_resource_dir)
 
     assert 'Test' in content
diff --git a/tests/k8s/test_k8s.py b/tests/k8s/test_k8s.py
index 3507c3984d678..6e80fc011f21b 100644
--- a/tests/k8s/test_k8s.py
+++ b/tests/k8s/test_k8s.py
@@ -1,4 +1,3 @@
-import time
 from http import HTTPStatus
 from pytest_kind import cluster
 
@@ -8,74 +7,23 @@
 import pytest
 import requests
 
-
 from jina import Flow
+from jina.peapods.pods.k8slib.kubernetes_tools import get_port_forward_contextmanager
 
 
-@pytest.fixture()
-def k8s_flow_with_needs(test_executor_image: str, executor_merger_image: str) -> Flow:
-    flow = (
-        Flow(name='test-flow', port_expose=9090, infrastructure='K8S', protocol='http')
-        .add(
-            name='segmenter',
-            uses=test_executor_image,
-        )
-        .add(
-            name='textencoder',
-            uses=test_executor_image,
-            needs='segmenter',
-        )
-        .add(
-            name='textstorage',
-            uses=test_executor_image,
-            needs='textencoder',
-        )
-        .add(
-            name='imageencoder',
-            uses=test_executor_image,
-            needs='segmenter',
-        )
-        .add(
-            name='imagestorage',
-            uses=test_executor_image,
-            needs='imageencoder',
-        )
-        .add(
-            name='merger',
-            uses=executor_merger_image,
-            needs=['imagestorage', 'textstorage'],
-        )
-    )
-    return flow
-
-
-def pull_images(images, cluster, logger):
-    # image pull anyways must be Never or IfNotPresent otherwise kubernetes will try to pull the image anyway
-    logger.debug(f'Loading docker image into kind cluster...')
-    for image in images:
-        cluster.needs_docker_image(image)
-    cluster.needs_docker_image('jinaai/jina:test-pip')
-    logger.debug(f'Done loading docker image into kind cluster...')
-
-
-def run_test(
-    images, cluster, flow, logger, expected_running_pods, endpoint, port_expose
-):
-    pull_images(images, cluster, logger)
-    start_flow(expected_running_pods, cluster, flow, logger)
-    resp = send_dummy_request(endpoint, cluster, flow, logger, port_expose=port_expose)
+def run_test(flow, logger, endpoint, port_expose):
+    with flow:
+        resp = send_dummy_request(endpoint, flow, logger, port_expose=port_expose)
+        flow.block()
     return resp
 
 
-def send_dummy_request(
-    endpoint, k8s_cluster_namespaced, k8s_flow_with_needs, logger, port_expose
-):
+def send_dummy_request(endpoint, flow, logger, port_expose):
     logger.debug(f'Starting port-forwarding to gateway service...')
-    with k8s_cluster_namespaced.port_forward(
-        'service/gateway', port_expose, port_expose, k8s_flow_with_needs.args.name
-    ) as _:
+    with get_port_forward_contextmanager(
+        namespace=flow.args.name, port_expose=port_expose
+    ):
         logger.debug(f'Port-forward running...')
-
         resp = requests.post(
             f'http://localhost:{port_expose}/{endpoint}',
             json={'data': [{} for _ in range(10)]},
@@ -83,43 +31,23 @@ def send_dummy_request(
     return resp
 
 
-def start_flow(expected_running_pods, k8s_cluster_namespaced, flow, logger):
-    logger.debug(f'Starting flow on kind cluster...')
-    flow.start()
-    logger.debug(f'Done starting flow on kind cluster...')
-    logger.debug(
-        f'Starting to wait for pods in kind cluster to reach "RUNNING" state...'
-    )
-    waiting = True
-    while waiting:
-        num_running_pods = len(
-            k8s_cluster_namespaced.list_ready_pods(namespace=flow.args.name)
-        )
-        if num_running_pods == expected_running_pods:
-            waiting = False
-        time.sleep(3)
-        logger.debug(
-            f'Still waiting for pods to reach running state '
-            f'(Current Status: {num_running_pods}/{expected_running_pods}).'
-        )
-
-        nodes_info = k8s_cluster_namespaced.get_node_info()
-        for n in nodes_info:
-            logger.debug(f'Node info {n}')
-
-
 @pytest.fixture()
 def k8s_flow_with_init_container(
     test_executor_image: str, executor_merger_image: str, dummy_dumper_image: str
 ) -> Flow:
     flow = Flow(
-        name='test-flow', port_expose=8080, infrastructure='K8S', protocol='http'
+        name='test-flow-with-init-container',
+        port_expose=9090,
+        infrastructure='K8S',
+        protocol='http',
+        timeout_ready=120000,
     ).add(
         name='test_executor',
         uses=test_executor_image,
         k8s_init_container_command=["python", "dump.py", "/shared/test_file.txt"],
         k8s_uses_init=dummy_dumper_image,
         k8s_mount_path='/shared',
+        timeout_ready=120000,
     )
     return flow
 
@@ -129,13 +57,18 @@ def k8s_flow_with_sharding(
     test_executor_image: str, executor_merger_image: str, dummy_dumper_image: str
 ) -> Flow:
     flow = Flow(
-        name='test-flow', port_expose=8080, infrastructure='K8S', protocol='http'
+        name='test-flow-with-sharding',
+        port_expose=9090,
+        infrastructure='K8S',
+        protocol='http',
+        timeout_ready=120000,
     ).add(
         name='test_executor',
-        shards=3,
+        shards=2,
         replicas=2,
         uses=test_executor_image,
         uses_after=executor_merger_image,
+        timeout_ready=360000,
     )
     return flow
 
@@ -143,20 +76,53 @@ def k8s_flow_with_sharding(
 @pytest.mark.timeout(3600)
 @pytest.mark.parametrize('k8s_connection_pool', [True, False])
 def test_flow_with_needs(
-    k8s_cluster_namespaced,
-    test_executor_image,
-    executor_merger_image,
-    k8s_flow_with_needs: Flow,
+    k8s_cluster,
+    test_executor_image: str,
+    executor_merger_image: str,
+    load_images_in_kind,
+    set_test_pip_version,
     logger,
     k8s_connection_pool: bool,
 ):
-    k8s_flow_with_needs.args.k8s_connection_pool = k8s_connection_pool
+    name = 'test-flow-with-needs'
+    if k8s_connection_pool:
+        name += '-pool'
+    flow = (
+        Flow(
+            name=name,
+            port_expose=9090,
+            infrastructure='K8S',
+            protocol='http',
+            timeout_ready=120000,
+            k8s_connection_pool=k8s_connection_pool,
+        )
+        .add(
+            name='segmenter',
+            uses=test_executor_image,
+            timeout_ready=120000,
+        )
+        .add(
+            name='textencoder',
+            uses=test_executor_image,
+            needs='segmenter',
+            timeout_ready=120000,
+        )
+        .add(
+            name='imageencoder',
+            uses=test_executor_image,
+            needs='segmenter',
+            timeout_ready=120000,
+        )
+        .add(
+            name='merger',
+            uses=executor_merger_image,
+            timeout_ready=120000,
+            needs=['imageencoder', 'textencoder'],
+        )
+    )
     resp = run_test(
-        [test_executor_image, executor_merger_image],
-        k8s_cluster_namespaced,
-        k8s_flow_with_needs,
+        flow,
         logger,
-        expected_running_pods=7,
         endpoint='index',
         port_expose=9090,
     )
@@ -165,8 +131,6 @@ def test_flow_with_needs(
         'segmenter',
         'imageencoder',
         'textencoder',
-        'imagestorage',
-        'textstorage',
     }
 
     assert resp.status_code == HTTPStatus.OK
@@ -178,20 +142,17 @@ def test_flow_with_needs(
 
 @pytest.mark.timeout(3600)
 def test_flow_with_init(
-    k8s_cluster_namespaced,
-    test_executor_image,
-    dummy_dumper_image: str,
+    k8s_cluster,
     k8s_flow_with_init_container: Flow,
+    load_images_in_kind,
+    set_test_pip_version,
     logger,
 ):
     resp = run_test(
-        [test_executor_image, dummy_dumper_image],
-        k8s_cluster_namespaced,
         k8s_flow_with_init_container,
         logger,
-        expected_running_pods=2,
         endpoint='search',
-        port_expose=8080,
+        port_expose=9090,
     )
 
     assert resp.status_code == HTTPStatus.OK
@@ -203,21 +164,17 @@ def test_flow_with_init(
 
 @pytest.mark.timeout(3600)
 def test_flow_with_sharding(
-    k8s_cluster_namespaced,
-    test_executor_image,
-    executor_merger_image,
+    k8s_cluster,
     k8s_flow_with_sharding: Flow,
+    load_images_in_kind,
+    set_test_pip_version,
     logger,
 ):
-
     resp = run_test(
-        [test_executor_image, executor_merger_image],
-        k8s_cluster_namespaced,
         k8s_flow_with_sharding,
         logger,
-        expected_running_pods=9,
         endpoint='index',
-        port_expose=8080,
+        port_expose=9090,
     )
 
     expected_traversed_executors = {
diff --git a/tests/k8s/test_k8s.sh b/tests/k8s/test_k8s.sh
deleted file mode 100755
index 54ce4471f13d2..0000000000000
--- a/tests/k8s/test_k8s.sh
+++ /dev/null
@@ -1,15 +0,0 @@
-set -ex
-
-docker build --build-arg PIP_TAG="[devel]" -f Dockerfiles/pip.Dockerfile -t jinaai/jina:test-pip .
-if [ "${PWD##*/}" != "jina" ]
-  then
-    echo "test_k8s.sh should only be run from the jina base directory"
-    exit 1
-fi
-
-pip install pytest
-pip install numpy
-export JINA_K8S_USE_TEST_PIP=True
-pytest -sv ./tests/k8s/test_k8s.py
-pytest -sv ./tests/k8s/test_custom_resource_dir.py
-
diff --git a/tests/unit/peapods/pods/k8slib/test_kubernetes_deployment.py b/tests/unit/peapods/pods/k8slib/test_kubernetes_deployment.py
index d141d162c5eed..19873d885563c 100644
--- a/tests/unit/peapods/pods/k8slib/test_kubernetes_deployment.py
+++ b/tests/unit/peapods/pods/k8slib/test_kubernetes_deployment.py
@@ -150,7 +150,7 @@ def test_init_container_args(pod_args):
     args = set_pod_parser().parse_args(pod_args)
     pod = K8sPod(args)
 
-    init_container = get_init_container_args(pod)
+    init_container = get_init_container_args(pod.args)
 
     if any(['--k8s-uses-init' in arg for arg in pod_args]):
         assert init_container == {
diff --git a/tests/unit/peapods/pods/k8slib/test_kubernetes_tools.py b/tests/unit/peapods/pods/k8slib/test_kubernetes_tools.py
index 9b7531121e89f..20707e2299eb3 100644
--- a/tests/unit/peapods/pods/k8slib/test_kubernetes_tools.py
+++ b/tests/unit/peapods/pods/k8slib/test_kubernetes_tools.py
@@ -6,18 +6,18 @@
 import kubernetes
 
 from jina.peapods.pods.k8slib.kubernetes_tools import create
-from jina.peapods.pods.k8slib.kubernetes_tools import __k8s_clients
+from jina.peapods.pods.k8slib.kubernetes_tools import _k8s_clients
 
 
 def test_lazy_load_k8s_client(monkeypatch):
     load_kube_config_mock = Mock()
     monkeypatch.setattr(kubernetes.config, 'load_kube_config', load_kube_config_mock)
-    attributes = ['k8s_client', 'v1', 'beta', 'networking_v1_beta1_api']
+    attributes = ['k8s_client', 'core_v1', 'beta', 'networking_v1_beta1_api', 'apps_v1']
     for attribute in attributes:
-        assert getattr(__k8s_clients, f'_K8SClients__{attribute}') is None
+        assert getattr(_k8s_clients, f'_{attribute}') is None
 
     for attribute in attributes:
-        assert getattr(__k8s_clients, attribute) is not None
+        assert getattr(_k8s_clients, attribute) is not None
 
 
 @pytest.mark.parametrize(
diff --git a/tests/unit/peapods/pods/test_k8s_pod.py b/tests/unit/peapods/pods/test_k8s_pod.py
index 7670b5ee4b186..f92e6bcbdfe29 100644
--- a/tests/unit/peapods/pods/test_k8s_pod.py
+++ b/tests/unit/peapods/pods/test_k8s_pod.py
@@ -27,10 +27,8 @@ def namespace_equal(
 
 
 @pytest.mark.parametrize('is_master', (True, False))
-def test_version(is_master, requests_mock, monkeypatch):
+def test_version(is_master, requests_mock):
     args = set_pod_parser().parse_args(['--name', 'test-pod'])
-    mock_create = Mock()
-    monkeypatch.setattr(kubernetes_tools, 'create', mock_create)
     if is_master:
         version = 'v2'
     else:
@@ -41,15 +39,10 @@ def test_version(is_master, requests_mock, monkeypatch):
         text='[{"name": "v1"}, {"name": "' + version + '"}]',
     )
     pod = K8sPod(args)
-
-    with pod:
-        assert (
-            mock_create.call_count == 5
-        )  # 3 because of namespace, service and deployment
-        if is_master:
-            assert pod.version == 'master'
-        else:
-            assert pod.version == jina.__version__
+    if is_master:
+        assert pod.version == 'master'
+    else:
+        assert pod.version == jina.__version__
 
 
 def test_dictionary_to_cli_param():
@@ -118,13 +111,13 @@ def test_parse_args_custom_executor(parallel: int):
             '2',
             [
                 {
-                    'name': 'test-pod_head',
+                    'name': 'test-pod-head',
                     'head_host': 'test-pod-head.ns.svc',
                 },
-                {'name': 'test-pod_0', 'head_host': 'test-pod-0.ns.svc'},
-                {'name': 'test-pod_1', 'head_host': 'test-pod-1.ns.svc'},
+                {'name': 'test-pod-0', 'head_host': 'test-pod-0.ns.svc'},
+                {'name': 'test-pod-1', 'head_host': 'test-pod-1.ns.svc'},
                 {
-                    'name': 'test-pod_tail',
+                    'name': 'test-pod-tail',
                     'head_host': 'test-pod-tail.ns.svc',
                 },
             ],
@@ -159,7 +152,6 @@ def get_k8s_pod(
     uses_after=None,
     port_expose=None,
 ):
-
     parameter_list = ['--name', pod_name, '--k8s-namespace', namespace]
     if parallel:
         parameter_list.extend(
@@ -192,6 +184,7 @@ def get_k8s_pod(
         )
     if uses_after:
         parameter_list.extend(['--uses-after', uses_after])
+    parameter_list.append('--noblock-on-start')
     parser = set_gateway_parser() if pod_name == 'gateway' else set_pod_parser()
     args = parser.parse_args(parameter_list)
     pod = K8sPod(args, needs)
@@ -202,11 +195,7 @@ def test_start_creates_namespace():
     ns = 'test'
     pod = get_k8s_pod('gateway', ns, port_expose=8085)
     kubernetes_deployment.deploy_service = Mock()
-    kubernetes_tools.create = Mock()
     pod.start()
-    kubernetes_tools.create.assert_called_once()
-    assert kubernetes_tools.create.call_args[0][0] == 'namespace'
-    assert kubernetes_tools.create.call_args[0][1] == {'name': ns}
     assert kubernetes_deployment.deploy_service.call_args[0][0] == 'gateway'
     assert kubernetes_deployment.deploy_service.call_args[1]['port_expose'] == 8085
 
@@ -217,7 +206,6 @@ def test_start_deploys_gateway():
 
     kubernetes_deployment.deploy_service = Mock()
     kubernetes_deployment.get_cli_params = Mock()
-    kubernetes_tools.create = Mock()
 
     pod = get_k8s_pod(pod_name, ns)
     pod.start()
@@ -239,9 +227,10 @@ def test_start_deploys_runtime():
     namespace = 'ns'
     pod = get_k8s_pod(pod_name, namespace)
 
-    pod._construct_runtime_container_args = Mock()
+    assert len(pod.k8s_deployments) > 0
+    for deployment in pod.k8s_deployments:
+        deployment._construct_runtime_container_args = Mock()
     kubernetes_deployment.deploy_service = Mock()
-    kubernetes_tools.create = Mock()
 
     pod.start()
 
@@ -256,12 +245,16 @@ def test_start_deploys_runtime():
     assert kwargs['init_container'] is None
     assert kwargs['custom_resource_dir'] is None
 
-    pod._construct_runtime_container_args.assert_called_once()
-    call_args = pod._construct_runtime_container_args.call_args[0]
-    assert call_args[0] == pod.deployment_args['deployments'][0]
-    assert call_args[1] == pod.args.uses
-    assert call_args[2] == kubernetes_deployment.dictionary_to_cli_param({'pea_id': 0})
-    assert call_args[3] == ''
+    assert len(pod.k8s_deployments) > 0
+    for i, deployment in enumerate(pod.k8s_deployments):
+        deployment._construct_runtime_container_args.assert_called_once()
+        call_args = deployment._construct_runtime_container_args.call_args[0]
+        assert call_args[0] == deployment.deployment_args
+        assert call_args[1] == pod.args.uses
+        assert call_args[2] == kubernetes_deployment.dictionary_to_cli_param(
+            {'pea_id': i}
+        )
+        assert call_args[3] == ''
 
 
 @pytest.mark.parametrize('parallel', [2, 3, 4])
@@ -271,7 +264,6 @@ def test_start_deploys_runtime_with_parallel(parallel: int):
 
     deploy_mock = Mock()
     kubernetes_deployment.deploy_service = deploy_mock
-    kubernetes_tools.create = Mock()
 
     pod.start()
 
@@ -311,7 +303,6 @@ def test_needs(needs, replicas, expected_calls, expected_executors):
 
     deploy_mock = Mock()
     kubernetes_deployment.deploy_service = deploy_mock
-    kubernetes_tools.create = Mock()
     pod.start()
     assert expected_calls == kubernetes_deployment.deploy_service.call_count
 
@@ -342,7 +333,6 @@ def test_uses_before_and_uses_after(
     )
     deploy_mock = Mock()
     kubernetes_deployment.deploy_service = deploy_mock
-    kubernetes_tools.create = Mock()
     pod.start()
     assert expected_calls == kubernetes_deployment.deploy_service.call_count
 


Code review comments:
diff_hunk: @@ -223,7 +223,7 @@ jobs:
           GITHUB_TOKEN: ${{ secrets.GITHUB_TOKEN }}
       - name: Test k8s
         run: |
-          pytest -v -s --suppress-no-test-exit-code --force-flaky --min-passes 1 --max-runs 5 --cov=jina --cov-report=xml ./tests/k8s/
+          pytest -v -s --suppress-no-test-exit-code --force-flaky --min-passes 1 --max-runs 5 --cov=jina --cov-report=xml ./tests/k8s/test_k8s.py::test_flow_with_sharding
path: .github/workflows/ci.yml
text: @author:
```suggestion
          pytest -v -s --suppress-no-test-exit-code --force-flaky --min-passes 1 --max-runs 5 --cov=jina --cov-report=xml ./tests/k8s
```