feat: Flow compatible with docarray v2 (jina-ai#5861)

Signed-off-by: Joan Fontanals Martinez <[email protected]>
Co-authored-by: Jina Dev Bot <[email protected]>
Co-authored-by: Alex Cureton-Griffiths <[email protected]>

Author: 3 people

.github/workflows/cd.yml

@@ -140,6 +140,7 @@ jobs:
       - name: Test
         id: test
         run: |
+          pytest --suppress-no-test-exit-code --force-flaky --min-passes 1 --max-runs 5 --cov=jina --cov-report=xml --timeout=600 -v -s --ignore-glob='tests/integration/hub_usage/dummyhub*' tests/unit/serve/runtimes/test_helper.py
           pytest --suppress-no-test-exit-code --force-flaky --min-passes 1 --max-runs 5 --cov=jina --cov-report=xml --timeout=600 -v -s --ignore-glob='tests/integration/hub_usage/dummyhub*' tests/integration/docarray_v2
           pytest --suppress-no-test-exit-code --force-flaky --min-passes 1 --max-runs 5 --cov=jina --cov-report=xml --timeout=600 -v -s --ignore-glob='tests/integration/hub_usage/dummyhub*' tests/integration/deployment_http_composite
           echo "flag it as jina for codeoverage"

.github/workflows/ci.yml

@@ -421,6 +421,7 @@ jobs:
       - name: Test
         id: test
         run: |
+          pytest --suppress-no-test-exit-code --force-flaky --min-passes 1 --max-runs 5 --cov=jina --cov-report=xml --timeout=600 -v -s --ignore-glob='tests/integration/hub_usage/dummyhub*' tests/unit/serve/runtimes/test_helper.py
           pytest --suppress-no-test-exit-code --force-flaky --min-passes 1 --max-runs 5 --cov=jina --cov-report=xml --timeout=600 -v -s --ignore-glob='tests/integration/hub_usage/dummyhub*' tests/integration/docarray_v2
           pytest --suppress-no-test-exit-code --force-flaky --min-passes 1 --max-runs 5 --cov=jina --cov-report=xml --timeout=600 -v -s --ignore-glob='tests/integration/hub_usage/dummyhub*' tests/integration/deployment_http_composite
           echo "flag it as jina for codeoverage"

docs/docarray-v2.md

@@ -92,9 +92,99 @@ If there is no `request_schema` and `response_schema`, the type hint is used to
 and `response_schema` will be used.
 
 
+## Serve one Executor in a Deployment
+
+Once you have defined the Executor with the New Executor API, you can easily serve and scale it as a Deployment with `gRPC`, `HTTP` or any combination of these
+protocols.
+
+
+```{code-block} python
+from jina import Deployment
+
+with Deployment(uses=MyExec, protocol='grpc', replicas=2) as dep:
+    dep.block()
+```
+
+
+## Chain Executors in Flow with different schemas
+
+With the new API, when building a Flow you should ensure that the Document types used as input of an Executor match the schema 
+of the output of its incoming previous Flow.
+
+For instance, this Flow will fail to start because the Document types are wrongly chained.
+
+````{tab} Valid Flow
+```{code-block} python
+from jina import Executor, requests, Flow
+from docarray import DocList, BaseDoc
+from docarray.typing import NdArray
+import numpy as np
+
+
+class SimpleStrDoc(BaseDoc):
+    text: str
+
+class TextWithEmbedding(SimpleStrDoc):
+    embedding: NdArray
+
+class TextEmbeddingExecutor(Executor):
+    @requests(on='/foo')
+    def foo(docs: DocList[SimpleStrDoc], **kwargs) -> DocList[TextWithEmbedding]
+        ret = DocList[TextWithEmbedding]()
+        for doc in docs:
+            ret.append(TextWithEmbedding(text=doc.text, embedding=np.ramdom.rand(10))
+        return ret
+
+class ProcessEmbedding(Executor):
+    @requests(on='/foo')
+    def foo(docs: DocList[TextWithEmbedding], **kwargs) -> DocList[TextWithEmbedding]
+        for doc in docs:
+            self.logger.info(f'Getting embedding with shape {doc.embedding.shape}')
+
+flow = Flow().add(uses=TextEmbeddingExecutor, name='embed').add(uses=ProcessEmbedding, name='process')
+with flow:
+    flow.block()
+```
+````
+````{tab} Invalid Flow
+```yaml
+from jina import Executor, requests, Flow
+from docarray import DocList, BaseDoc
+from docarray.typing import NdArray
+import numpy as np
+
+
+class SimpleStrDoc(BaseDoc):
+    text: str
+
+class TextWithEmbedding(SimpleStrDoc):
+    embedding: NdArray
+
+class TextEmbeddingExecutor(Executor):
+    @requests(on='/foo')
+    def foo(docs: DocList[SimpleStrDoc], **kwargs) -> DocList[TextWithEmbedding]
+        ret = DocList[TextWithEmbedding]()
+        for doc in docs:
+            ret.append(TextWithEmbedding(text=doc.text, embedding=np.ramdom.rand(10))
+        return ret
+
+class ProcessText(Executor):
+    @requests(on='/foo')
+    def foo(docs: DocList[SimpleStrDoc], **kwargs) -> DocList[TextWithEmbedding]
+        for doc in docs:
+            self.logger.info(f'Getting embedding with type {doc.text}')
+
+# This Flow will fail to start because the input type of "process" does not match the output type of "embed"
+flow = Flow().add(uses=TextEmbeddingExecutor, name='embed').add(uses=ProcessText, name='process')
+with flow:
+    flow.block()
+```
+````
+
+
 ## Client API
 
-In the client, you similarly specify the schema that you expect the Flow to return. You can pass the return type by using the `return_type` parameter in the `client.post` method:
+Similarly, In the client, you specify the schema that you expect the Deployment or Flow to return. You can pass the return type by using the `return_type` parameter in the `client.post` method:
 
 ```{code-block} python
 ---
@@ -117,20 +207,13 @@ with Deployment(uses=MyExec) as dep:
 
 Jina is working to offer full compatibility with the new DocArray version.
 
-At present, these features are supported if you use the APIs described in the previous sections.
-
-- All the features offered by {ref}`Deployment <deployment>` where a single Executor is served and {ref}`scaled <scale-out>`. This includes both
-HTTP and gRPC protocols, and both of them at the same type.
-
-- When combining multiple Deployments in a pipeline using a {ref}`Flow <flow-cookbook>`, there are currently several limitations:
+However, there are currently some limitations to consider.
 
-    - Only gRPC protocol is supported.
-    - Only linear Flows are supported, no topologies using bifurcations can be used at the moment.
 
 ````{admonition} Note
 :class: note
 
-With DocArray 0.30 support, Jina introduced the concept of input/output schema at the Executor level. In order to chain multiple Executor into a Flow you always need to make sure that the output schema of an Executor is the same as the Input of the Executor that follows him in the Flow
+With DocArray 0.30 support, Jina introduced the concept of input/output schema at the Executor level. To chain multiple Executors into a Flow you need to ensure that the output schema of an Executor is the same as the input of the Executor that follows it in the Flow
 ```
 
 ````{admonition} Note

jina/clients/request/helper.py

@@ -6,7 +6,7 @@
 from jina.types.request.data import DataRequest
 
 if docarray_v2:
-    from docarray import DocList
+    from docarray import DocList, BaseDoc
 
 
 def _new_data_request_from_batch(
@@ -80,7 +80,11 @@ def _add_docs(req: DataRequest, batch, data_type: DataInputType) -> None:
     if not docarray_v2:
         da = DocumentArray([])
     else:
-        da = DocList[batch[0].__class__]()
+        if len(batch) > 0:
+            da = DocList[batch[0].__class__]()
+        else:
+            da = DocList[BaseDoc]()
+
     for content in batch:
         d, data_type = _new_doc_from_data(content, data_type)
         da.append(d)

jina/serve/consensus/run.go

@@ -359,7 +359,6 @@ func run(self *C.PyObject, args *C.PyObject, kwargs *C.PyObject) *C.PyObject {
 
 //export add_voter
 func add_voter(self *C.PyObject, args *C.PyObject) *C.PyObject {
-    //TODO: Instantiate new logger based on JINA_LOG_LEVEL
     logLevel := os.Getenv("JINA_LOG_LEVEL")
     if logLevel == "" {
         logLevel = "INFO"

jina/serve/networking/__init__.py

@@ -203,15 +203,15 @@ def send_discover_endpoint(
             shard_id: Optional[int] = None,
             timeout: Optional[float] = None,
             retries: Optional[int] = -1,
-    ) -> Optional[asyncio.Task]:
+    ):
         """Sends a discover Endpoint call to target.
 
         :param deployment: name of the Jina deployment to send the request to
         :param head: If True it is send to the head, otherwise to the worker pods
         :param shard_id: Send to a specific shard of the deployment, ignored for polling ALL
         :param timeout: timeout for sending the requests
         :param retries: number of retries per gRPC call. If <0 it defaults to max(3, num_replicas)
-        :return: asyncio.Task items to send call
+        :return: coroutine items to send call
         """
         connection_list = self._connections.get_replicas(
             deployment, head, shard_id, True
@@ -378,12 +378,9 @@ async def _handle_aiorpcerror(
                 details=error.details(),
             )
         else:
-            if error.code() == grpc.StatusCode.UNAVAILABLE and 'not the leader' in error.details():
-                self._logger.debug(f'RAFT node of {current_deployment} is not the leader. Trying next replica, if available.')
-            else:
-                self._logger.debug(
-                    f'gRPC call to deployment {current_deployment} failed with error {format_grpc_error(error)}, for retry attempt {retry_i + 1}/{total_num_tries - 1}.'
-                    f' Trying next replica, if available.'
+            if connection_list:
+                await connection_list.reset_connection(
+                    current_address, current_deployment
                 )
             return None
 
@@ -460,11 +457,10 @@ def _send_discover_endpoint(
             connection_list: _ReplicaList,
             timeout: Optional[float] = None,
             retries: Optional[int] = -1,
-    ) -> asyncio.Task:
+    ):
         # this wraps the awaitable object from grpc as a coroutine so it can be used as a task
         # the grpc call function is not a coroutine but some _AioCall
-        async def task_wrapper():
-
+        async def task_coroutine():
             tried_addresses = set()
             if retries is None or retries < 0:
                 total_num_tries = (
@@ -500,7 +496,7 @@ async def task_wrapper():
                 except AttributeError:
                     return default_endpoints_proto, None
 
-        return asyncio.create_task(task_wrapper())
+        return task_coroutine()
 
     async def warmup(
             self,
@@ -557,7 +553,6 @@ async def task_wrapper(target_warmup_responses, stub):
                         for task in tasks:
                             task.cancel()
                     raise
-
         except Exception as ex:
             self._logger.error(f'error with warmup up task: {ex}')
             return
@@ -568,5 +563,4 @@ def _get_all_replicas(self, deployment):
         replica_set.add(
             self._connections.get_replicas(deployment=deployment, head=True)
         )
-
         return set(filter(None, replica_set))

jina/serve/networking/replica_list.py

@@ -23,15 +23,15 @@ class _ReplicaList:
     """
 
     def __init__(
-        self,
-        metrics: _NetworkingMetrics,
-        histograms: _NetworkingHistograms,
-        logger,
-        runtime_name: str,
-        aio_tracing_client_interceptors: Optional[Sequence['ClientInterceptor']] = None,
-        tracing_client_interceptor: Optional['OpenTelemetryClientInterceptor'] = None,
-        deployment_name: str = '',
-        channel_options: Optional[Union[list, Dict[str, Any]]] = None,
+            self,
+            metrics: _NetworkingMetrics,
+            histograms: _NetworkingHistograms,
+            logger,
+            runtime_name: str,
+            aio_tracing_client_interceptors: Optional[Sequence['ClientInterceptor']] = None,
+            tracing_client_interceptor: Optional['OpenTelemetryClientInterceptor'] = None,
+            deployment_name: str = '',
+            channel_options: Optional[Union[list, Dict[str, Any]]] = None,
     ):
         self.runtime_name = runtime_name
         self._connections = []
@@ -59,20 +59,21 @@ async def reset_connection(self, address: str, deployment_name: str):
         :param deployment_name: Target deployment of this connection
         """
         self._logger.debug(f'resetting connection for {deployment_name} to {address}')
-
+        parsed_address = urlparse(address)
+        resolved_address = parsed_address.netloc if parsed_address.netloc else address
         if (
-            address in self._address_to_connection_idx
-            and self._address_to_connection_idx[address] is not None
+                resolved_address in self._address_to_connection_idx
+                and self._address_to_connection_idx[resolved_address] is not None
         ):
             # remove connection:
             # in contrast to remove_connection(), we don't 'shorten' the data structures below, instead
             # update the data structure with the new connection and let the old connection be colleced by
             # the GC
-            id_to_reset = self._address_to_connection_idx[address]
+            id_to_reset = self._address_to_connection_idx[resolved_address]
             # re-add connection:
-            self._address_to_connection_idx[address] = id_to_reset
+            self._address_to_connection_idx[resolved_address] = id_to_reset
             stubs, channel = self._create_connection(address, deployment_name)
-            self._address_to_channel[address] = channel
+            self._address_to_channel[resolved_address] = channel
             self._connections[id_to_reset] = stubs
 
     def add_connection(self, address: str, deployment_name: str):
@@ -81,10 +82,13 @@ def add_connection(self, address: str, deployment_name: str):
         :param address: Target address of this connection
         :param deployment_name: Target deployment of this connection
         """
-        if address not in self._address_to_connection_idx:
-            self._address_to_connection_idx[address] = len(self._connections)
+        parsed_address = urlparse(address)
+        resolved_address = parsed_address.netloc if parsed_address.netloc else address
+
+        if resolved_address not in self._address_to_connection_idx:
+            self._address_to_connection_idx[resolved_address] = len(self._connections)
             stubs, channel = self._create_connection(address, deployment_name)
-            self._address_to_channel[address] = channel
+            self._address_to_channel[resolved_address] = channel
             self._connections.append(stubs)
             # create a new set of stubs and channels for warmup to avoid
             # loosing channel during remove_connection or reset_connection
@@ -103,20 +107,23 @@ async def remove_connection(self, address: str):
 
         :param address: Remove connection for this address
         """
-        if address in self._address_to_connection_idx:
+        parsed_address = urlparse(address)
+        resolved_address = parsed_address.netloc if parsed_address.netloc else address
+        if resolved_address in self._address_to_connection_idx:
             self._rr_counter = (
                 self._rr_counter % (len(self._connections) - 1)
                 if (len(self._connections) - 1)
                 else 0
             )
-            idx_to_delete = self._address_to_connection_idx.pop(address)
+            idx_to_delete = self._address_to_connection_idx.pop(resolved_address)
             self._connections.pop(idx_to_delete)
             # update the address/idx mapping
-            for address in self._address_to_connection_idx:
-                if self._address_to_connection_idx[address] > idx_to_delete:
-                    self._address_to_connection_idx[address] -= 1
+            for a in self._address_to_connection_idx:
+                if self._address_to_connection_idx[a] > idx_to_delete:
+                    self._address_to_connection_idx[a] -= 1
 
     def _create_connection(self, address, deployment_name: str):
+        self._logger.debug(f'create_connection connection for {deployment_name} to {address}')
         parsed_address = urlparse(address)
         address = parsed_address.netloc if parsed_address.netloc else address
         use_tls = parsed_address.scheme in TLS_PROTOCOL_SCHEMES
@@ -185,7 +192,9 @@ def has_connection(self, address: str) -> bool:
         :param address: The address to check
         :returns: True if a connection for the ip exists in the list
         """
-        return address in self._address_to_connection_idx
+        parsed_address = urlparse(address)
+        resolved_address = parsed_address.netloc if parsed_address.netloc else address
+        return resolved_address in self._address_to_connection_idx
 
     def has_connections(self) -> bool:
         """