fix: fix isse in head missmatch endpoint (jina-ai#5904)

Author: JoanFM

jina/serve/runtimes/head/request_handling.py

@@ -9,6 +9,7 @@
 import grpc
 
 from jina.enums import PollingType
+from jina.constants import __default_endpoint__
 from jina.excepts import InternalNetworkError
 from jina.helper import get_full_version
 from jina.proto import jina_pb2
@@ -289,7 +290,10 @@ async def _handle_data_request(
         response_request = worker_results[0]
         found = False
         if docarray_v2:
-            model = self._pydantic_models_by_endpoint[endpoint]['output']
+            check_endpoint = endpoint
+            if endpoint not in self._pydantic_models_by_endpoint:
+                check_endpoint = __default_endpoint__
+            model = self._pydantic_models_by_endpoint[check_endpoint]['output']
         for i, worker_result in enumerate(worker_results):
             if docarray_v2:
                 worker_result.document_array_cls = DocList[model]

tests/integration/docarray_v2/test_v2.py

@@ -987,3 +987,30 @@ def search(self, docs: DocList[TextDocWithId], **kwargs) -> DocList[ResultTestDo
                 assert len(r.matches) == 6
                 for match in r.matches:
                     assert 'ID' in match.text
+
+
+def test_issue_shards_missmatch_endpoint():
+
+    class MyDoc(BaseDoc):
+        text: str
+        embedding: NdArray[128]
+
+    class MyDocWithMatchesAndScores(MyDoc):
+        matches: DocList[MyDoc]
+        scores: List[float]
+
+    class MyExec(Executor):
+
+        @requests
+        def foo(self, docs: DocList[MyDoc], **kwargs) -> DocList[MyDocWithMatchesAndScores]:
+            res = DocList[MyDocWithMatchesAndScores]()
+            for doc in docs:
+                new_doc = MyDocWithMatchesAndScores(text=doc.text, embedding=doc.embedding, matches=docs,
+                                                    scores=[1.0 for _ in docs])
+                res.append(new_doc)
+            return res
+
+    d = Deployment(uses=MyExec, shards=2)
+    with d:
+        res = d.post(on='/', inputs=DocList[MyDoc]([MyDoc(text='hey ha', embedding=np.random.rand(128))]))
+        assert len(res) == 1