Repository: jina-ai/serve
Pull request: feat: start portforwarding when flow.post in K8s
This pr enables to run k8s flows via context manager.

# Example:
## Before:

```
    Flow(
        name='test-flow', port_expose=8080, infrastructure='K8S'
    ).start()
    
    # ... setup port forward ...
    with k8s_cluster_namespaced.port_forward(
        resp = requests.post(
            f'http://localhost:{port_expose}/{endpoint}',
            json={'data': [{} for _ in range(10)]},
        )

```

## After:
```
    with Flow(
        name='test-flow', port_expose=8080, infrastructure='K8S'
    ):
        flow.post(endpoint, [Document() for _ in range(10)])
```

In order to archive this, it uses port-forwarding during `flow.post`. 
For now, each call of `flow.post` creates and terminates port-forwarding. This means, calling `flow.post` multiple times creates an overhead, which can be tackled in a future pr.

If you are fine with the draft, I will add the docs as well in the same pr.


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

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

Commit message: refactor: probe

Patch to review:
diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml
index bd4845bd47660..ff1b071e5273c 100644
--- a/.github/workflows/ci.yml
+++ b/.github/workflows/ci.yml
@@ -322,7 +322,7 @@ jobs:
   # just for blocking the merge until all parallel core-test are successful
   success-all-test:
     runs-on: ubuntu-latest
-    needs: [core-test, hub-test, docker-image-test, check-docstring, check-black, code-injection]
+    needs: [core-test, hub-test, k8s-test, docker-image-test, check-docstring, check-black, code-injection]
     if: always()
     steps:
       - uses: technote-space/workflow-conclusion-action@v2
diff --git a/extra-requirements.txt b/extra-requirements.txt
index 051f311829886..7c4de5fe9e369 100644
--- a/extra-requirements.txt
+++ b/extra-requirements.txt
@@ -72,7 +72,7 @@ jsonschema:                 cicd
 dgl:                        cicd
 black==20.8b1:              test
 matplotlib:                 devel
+portforward>=0.2.4:         standard, devel
 kubernetes>=18.20.0:        perf, standard, devel
 pytest-kind==21.1.3:        test
 pytest-lazy-fixture:        test
-portforward>=0.2.4:         test
\ No newline at end of file
diff --git a/jina/clients/mixin.py b/jina/clients/mixin.py
index 9d2c4cbd3d2d9..f0fc441402517 100644
--- a/jina/clients/mixin.py
+++ b/jina/clients/mixin.py
@@ -1,8 +1,11 @@
+from contextlib import nullcontext
 from functools import partialmethod
 from typing import Optional, Dict, List, AsyncGenerator
 
 from .base import CallbackFnType, InputType
+from ..enums import InfrastructureType
 from ..helper import run_async
+from ..peapods.pods.k8slib import kubernetes_tools
 from ..types.request import Response
 
 
@@ -57,18 +60,28 @@ async def _get_results(*args, **kwargs):
             if return_results:
                 return result
 
-        return run_async(
-            _get_results,
-            inputs=inputs,
-            on_done=on_done,
-            on_error=on_error,
-            on_always=on_always,
-            exec_endpoint=on,
-            target_peapod=target_peapod,
-            parameters=parameters,
-            request_size=request_size,
-            **kwargs,
-        )
+        if (
+            hasattr(self.args, 'infrastructure')
+            and self.args.infrastructure == InfrastructureType.K8S
+        ):
+            context_mgr = kubernetes_tools.get_port_forward_contextmanager(
+                self.args.name, self.port_expose
+            )
+        else:
+            context_mgr = nullcontext()
+        with context_mgr:
+            return run_async(
+                _get_results,
+                inputs=inputs,
+                on_done=on_done,
+                on_error=on_error,
+                on_always=on_always,
+                exec_endpoint=on,
+                target_peapod=target_peapod,
+                parameters=parameters,
+                request_size=request_size,
+                **kwargs,
+            )
 
     # ONLY CRUD, for other request please use `.post`
     index = partialmethod(post, '/index')
diff --git a/jina/peapods/pods/k8s.py b/jina/peapods/pods/k8s.py
index b51781809cdaf..4b4121e43ebcf 100644
--- a/jina/peapods/pods/k8s.py
+++ b/jina/peapods/pods/k8s.py
@@ -311,8 +311,12 @@ def port_expose(self) -> int:
 
     @property
     def host(self) -> str:
-        """Not implemented"""
-        raise NotImplementedError
+        """Currently, when deploying on Kubernetes, Jina does not expose a public host.
+        Instead Jina sends requests via port-forward and runs the requests against localhost.
+
+        :return: localhost
+        """
+        return 'localhost'
 
     def start(self) -> 'K8sPod':
         """Deploy the kubernetes pods via k8s Deployment and k8s Service.
diff --git a/jina/peapods/pods/k8slib/kubernetes_tools.py b/jina/peapods/pods/k8slib/kubernetes_tools.py
index e0cfc043e6eaf..680588cc63861 100644
--- a/jina/peapods/pods/k8slib/kubernetes_tools.py
+++ b/jina/peapods/pods/k8slib/kubernetes_tools.py
@@ -1,7 +1,7 @@
+import json
 import os
 import tempfile
-import json
-from typing import Dict, Optional, TYPE_CHECKING
+from typing import Dict, Optional, Generator
 
 from ....importer import ImportExtensions
 from ....logging.logger import JinaLogger
@@ -12,9 +12,6 @@
     cur_dir, '..', '..', '..', 'resources', 'k8s', 'template'
 )
 
-if TYPE_CHECKING:
-    from contextlib import GeneratorContextManager
-
 
 class K8sClients:
     """
@@ -167,7 +164,7 @@ def get_port_forward_contextmanager(
     namespace: str,
     port_expose: int,
     config_path: str = None,
-) -> 'GeneratorContextManager':
+) -> Generator[None, None, None]:
     """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
diff --git a/jina/resources/extra-requirements.txt b/jina/resources/extra-requirements.txt
index 051f311829886..7c4de5fe9e369 100644
--- a/jina/resources/extra-requirements.txt
+++ b/jina/resources/extra-requirements.txt
@@ -72,7 +72,7 @@ jsonschema:                 cicd
 dgl:                        cicd
 black==20.8b1:              test
 matplotlib:                 devel
+portforward>=0.2.4:         standard, devel
 kubernetes>=18.20.0:        perf, standard, devel
 pytest-kind==21.1.3:        test
 pytest-lazy-fixture:        test
-portforward>=0.2.4:         test
\ No newline at end of file
diff --git a/jina/resources/k8s/template/deployment-init.yml b/jina/resources/k8s/template/deployment-init.yml
index fa93c8f31a4dd..3cc88c8a1c1f7 100644
--- a/jina/resources/k8s/template/deployment-init.yml
+++ b/jina/resources/k8s/template/deployment-init.yml
@@ -30,6 +30,13 @@ spec:
       containers:
       - name: executor
         image: {image}
+#        resources:
+#          requests:
+#            memory: "256Mi"
+#            cpu: "250m"
+#          limits:
+#            memory: "256Mi"
+#            cpu: "250m"
         command: {command}
         args: {args}
         ports:
@@ -42,18 +49,13 @@ spec:
           value: debug
         - name: PYTHONUNBUFFERED
           value: "1"
-        - name: JINA_EXECUTOR_CONTROL_PORT
-          value: "8081" # TODO there should be a better way to pass this and align with port_ctrl
         - name: WORKER_CLASS
           value: "uvicorn.workers.UvicornH11Worker"
         volumeMounts:
           - name: shared
             mountPath: {mount-path}
         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)'
+          tcpSocket:
+            port: {port_in}
           initialDelaySeconds: 5
-          periodSeconds: 5
+          periodSeconds: 10
diff --git a/jina/resources/k8s/template/deployment.yml b/jina/resources/k8s/template/deployment.yml
index 602b15be8d323..73a5b0bf4140a 100644
--- a/jina/resources/k8s/template/deployment.yml
+++ b/jina/resources/k8s/template/deployment.yml
@@ -17,6 +17,14 @@ spec:
       containers:
       - name: executor
         image: {image}
+#        resources:
+#          requests:
+#            memory: "256Mi"
+#            cpu: "250m"
+#          limits:
+#            memory: "256Mi"
+#            cpu: "250m"
+
         imagePullPolicy: {pull_policy}
         command: {command}
         args: {args}
@@ -26,13 +34,11 @@ spec:
           - 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)'
+          tcpSocket:
+            port: {port_in}
           initialDelaySeconds: 5
-          periodSeconds: 5
+          periodSeconds: 10
+
 #        livenessProbe:
 #          exec:
 #            command:
@@ -48,5 +54,4 @@ spec:
             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/k8s/dummy-dumper/Dockerfile b/tests/k8s/dummy-dumper/Dockerfile
index 1280992564013..b740c24b36ae4 100644
--- a/tests/k8s/dummy-dumper/Dockerfile
+++ b/tests/k8s/dummy-dumper/Dockerfile
@@ -1,4 +1,3 @@
-# TODO use fixed jina version for deterministic execution
 FROM jinaai/jina:test-pip
 
 # setup the workspace
diff --git a/tests/k8s/test_k8s.py b/tests/k8s/test_k8s.py
index 84b49581eb888..9710203b4a23a 100644
--- a/tests/k8s/test_k8s.py
+++ b/tests/k8s/test_k8s.py
@@ -1,31 +1,20 @@
-from http import HTTPStatus
 from pytest_kind import cluster
 
 # kind version has to be bumped to v0.11.1 since pytest-kind is just using v0.10.0 which does not work on ubuntu in ci
 # TODO don't use pytest-kind anymore
 cluster.KIND_VERSION = 'v0.11.1'
 import pytest
-import requests
 
-from jina import Flow
-from jina.peapods.pods.k8slib.kubernetes_tools import get_port_forward_contextmanager
+from jina import Flow, Document
 
 
-def run_test(flow, logger, endpoint, port_expose):
+def run_test(flow, endpoint, port_expose):
     with flow:
-        resp = send_dummy_request(endpoint, flow, logger, port_expose=port_expose)
-    return resp
-
-
-def send_dummy_request(endpoint, flow, logger, port_expose):
-    logger.debug(f'Starting port-forwarding to gateway service...')
-    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)]},
+        resp = flow.post(
+            endpoint,
+            [Document() for _ in range(10)],
+            return_results=True,
+            port_expose=port_expose,
         )
     return resp
 
@@ -121,8 +110,7 @@ def test_flow_with_needs(
     )
     resp = run_test(
         flow,
-        logger,
-        endpoint='index',
+        endpoint='/index',
         port_expose=9090,
     )
 
@@ -132,11 +120,10 @@ def test_flow_with_needs(
         'textencoder',
     }
 
-    assert resp.status_code == HTTPStatus.OK
-    docs = resp.json()['data']['docs']
+    docs = resp[0].docs
     assert len(docs) == 10
     for doc in docs:
-        assert set(doc['tags']['traversed-executors']) == expected_traversed_executors
+        assert set(doc.tags['traversed-executors']) == expected_traversed_executors
 
 
 @pytest.mark.timeout(3600)
@@ -149,16 +136,14 @@ def test_flow_with_init(
 ):
     resp = run_test(
         k8s_flow_with_init_container,
-        logger,
-        endpoint='search',
+        endpoint='/search',
         port_expose=9090,
     )
 
-    assert resp.status_code == HTTPStatus.OK
-    docs = resp.json()['data']['docs']
+    docs = resp[0].docs
     assert len(docs) == 10
     for doc in docs:
-        assert doc['tags']['file'] == ['1\n', '2\n', '3']
+        assert doc.tags['file'] == ['1\n', '2\n', '3']
 
 
 @pytest.mark.timeout(3600)
@@ -171,8 +156,7 @@ def test_flow_with_sharding(
 ):
     resp = run_test(
         k8s_flow_with_sharding,
-        logger,
-        endpoint='index',
+        endpoint='/index',
         port_expose=9090,
     )
 
@@ -180,8 +164,7 @@ def test_flow_with_sharding(
         'test_executor',
     }
 
-    assert resp.status_code == HTTPStatus.OK
-    docs = resp.json()['data']['docs']
+    docs = resp[0].docs
     assert len(docs) == 10
     for doc in docs:
-        assert set(doc['tags']['traversed-executors']) == expected_traversed_executors
+        assert set(doc.tags['traversed-executors']) == expected_traversed_executors


Code review comments:
diff_hunk: @@ -17,6 +17,14 @@ spec:
       containers:
       - name: executor
         image: {image}
+#        resources:
+#          requests:
path: jina/resources/k8s/template/deployment.yml
text: @user1:
remember to remove these