From ecdf2b645e71628491c7d4ed74fec7e713464c5b Mon Sep 17 00:00:00 2001 From: Joan Martinez Date: Thu, 7 Mar 2024 09:56:47 +0100 Subject: [PATCH] feat: adapt pydantic v2 --- jina/_docarray.py | 6 ++ .../runtimes/gateway/graph/topology_graph.py | 14 +++-- jina/serve/runtimes/head/request_handling.py | 11 ++-- jina/serve/runtimes/helper.py | 4 +- .../serve/runtimes/worker/request_handling.py | 11 ++-- tests/integration/docarray_v2/test_v2.py | 35 ++++++------ tests/unit/serve/runtimes/test_helper.py | 55 ++++++++++--------- 7 files changed, 80 insertions(+), 56 deletions(-) diff --git a/jina/_docarray.py b/jina/_docarray.py index b9cdfe14aa3fd..96ae3dd1d352d 100644 --- a/jina/_docarray.py +++ b/jina/_docarray.py @@ -8,3 +8,9 @@ from docarray import Document, DocumentArray docarray_v2 = False + + +import pydantic + +is_pydantic_v2 = pydantic.__version__.startswith('2.') + diff --git a/jina/serve/runtimes/gateway/graph/topology_graph.py b/jina/serve/runtimes/gateway/graph/topology_graph.py index 89ad937c4698a..2a6c01ea9b633 100644 --- a/jina/serve/runtimes/gateway/graph/topology_graph.py +++ b/jina/serve/runtimes/gateway/graph/topology_graph.py @@ -7,7 +7,7 @@ import grpc.aio -from jina._docarray import DocumentArray, docarray_v2 +from jina._docarray import DocumentArray, docarray_v2, is_pydantic_v2 from jina.constants import __default_endpoint__ from jina.excepts import InternalNetworkError from jina.logging.logger import JinaLogger @@ -20,7 +20,11 @@ from docarray import DocList from docarray.documents.legacy import LegacyDocument - from jina.serve.runtimes.helper import _create_pydantic_model_from_schema + if not is_pydantic_v2: + from jina.serve.runtimes.helper import _create_pydantic_model_from_schema as create_base_doc_from_schema + else: + from docarray.utils.create_dynamic_doc_class import create_base_doc_from_schema + legacy_doc_schema = LegacyDocument.schema() @@ -239,7 +243,7 @@ async def task(): input_model = LegacyDocument else: input_model = ( - _create_pydantic_model_from_schema( + create_base_doc_from_schema( input_model_schema, input_model_name, models_created_by_name, @@ -269,7 +273,7 @@ async def task(): output_model = LegacyDocument else: output_model = ( - _create_pydantic_model_from_schema( + create_base_doc_from_schema( output_model_schema, output_model_name, models_created_by_name, @@ -306,7 +310,7 @@ async def task(): from pydantic import BaseModel parameters_model = ( - _create_pydantic_model_from_schema( + create_base_doc_from_schema( parameters_model_schema, parameters_model_name, models_created_by_name, diff --git a/jina/serve/runtimes/head/request_handling.py b/jina/serve/runtimes/head/request_handling.py index 6891c68c02d6a..417c7a865ac6d 100644 --- a/jina/serve/runtimes/head/request_handling.py +++ b/jina/serve/runtimes/head/request_handling.py @@ -16,10 +16,13 @@ from jina.serve.runtimes.monitoring import MonitoringRequestMixin from jina.serve.runtimes.worker.request_handling import WorkerRequestHandler from jina.types.request.data import DataRequest, Response -from jina._docarray import docarray_v2 +from jina._docarray import docarray_v2, is_pydantic_v2 if docarray_v2: - from jina.serve.runtimes.helper import _create_pydantic_model_from_schema + if not is_pydantic_v2: + from jina.serve.runtimes.helper import _create_pydantic_model_from_schema as create_base_doc_from_schema + else: + from docarray.utils.create_dynamic_doc_class import create_base_doc_from_schema from docarray import DocList from docarray.base_doc.any_doc import AnyDoc @@ -359,7 +362,7 @@ async def task(): LegacyDocument ) elif input_model_name not in models_created_by_name: - input_model = _create_pydantic_model_from_schema( + input_model = create_base_doc_from_schema( input_model_schema, input_model_name, {} ) models_created_by_name[input_model_name] = input_model @@ -369,7 +372,7 @@ async def task(): LegacyDocument ) elif output_model_name not in models_created_by_name: - output_model = _create_pydantic_model_from_schema( + output_model = create_base_doc_from_schema( output_model_schema, output_model_name, {} ) models_created_by_name[output_model_name] = output_model diff --git a/jina/serve/runtimes/helper.py b/jina/serve/runtimes/helper.py index 70bb75a485c1b..6444f8fa5a03a 100644 --- a/jina/serve/runtimes/helper.py +++ b/jina/serve/runtimes/helper.py @@ -1,7 +1,7 @@ import copy from typing import Any, Dict, List, Optional, Tuple, Union -from jina._docarray import docarray_v2 +from jina._docarray import docarray_v2, is_pydantic_v2 _SPECIFIC_EXECUTOR_SEPARATOR = '__' @@ -79,7 +79,7 @@ def _parse_specific_params(parameters: Dict, executor_name: str): return parsed_params -if docarray_v2: +if docarray_v2 and not is_pydantic_v2: from docarray import BaseDoc, DocList from docarray.typing import AnyTensor from pydantic import create_model diff --git a/jina/serve/runtimes/worker/request_handling.py b/jina/serve/runtimes/worker/request_handling.py index 2e095cb26da50..65472cd6d406f 100644 --- a/jina/serve/runtimes/worker/request_handling.py +++ b/jina/serve/runtimes/worker/request_handling.py @@ -20,7 +20,7 @@ from google.protobuf.struct_pb2 import Struct -from jina._docarray import DocumentArray, docarray_v2 +from jina._docarray import DocumentArray, docarray_v2, is_pydantic_v2 from jina.constants import __default_endpoint__ from jina.excepts import BadConfigSource, RuntimeTerminated from jina.helper import get_full_version @@ -1013,21 +1013,24 @@ async def endpoint_discovery(self, empty, context) -> jina_pb2.EndpointsProto: if docarray_v2: from docarray.documents.legacy import LegacyDocument - from jina.serve.runtimes.helper import _create_aux_model_doc_list_to_list + if not is_pydantic_v2: + from jina.serve.runtimes.helper import _create_aux_model_doc_list_to_list as create_pure_python_type_model + else: + from docarray.utils.create_dynamic_doc_class import create_pure_python_type_model legacy_doc_schema = LegacyDocument.schema() for endpoint_name, inner_dict in schemas.items(): if inner_dict['input']['model'].schema() == legacy_doc_schema: inner_dict['input']['model'] = legacy_doc_schema else: - inner_dict['input']['model'] = _create_aux_model_doc_list_to_list( + inner_dict['input']['model'] = create_pure_python_type_model( inner_dict['input']['model'] ).schema() if inner_dict['output']['model'].schema() == legacy_doc_schema: inner_dict['output']['model'] = legacy_doc_schema else: - inner_dict['output']['model'] = _create_aux_model_doc_list_to_list( + inner_dict['output']['model'] = create_pure_python_type_model( inner_dict['output']['model'] ).schema() diff --git a/tests/integration/docarray_v2/test_v2.py b/tests/integration/docarray_v2/test_v2.py index eebbafd1a572f..7643be4951af8 100644 --- a/tests/integration/docarray_v2/test_v2.py +++ b/tests/integration/docarray_v2/test_v2.py @@ -14,6 +14,7 @@ import numpy as np import pytest +from jina._docarray import is_pydantic_v2 from docarray import BaseDoc, DocList from docarray.documents import ImageDoc, TextDoc from docarray.documents.legacy import LegacyDocument @@ -302,10 +303,11 @@ def bar(self, docs: DocList[Output1], **kwargs) -> DocList[Output2]: from jina.proto import jina_pb2 from jina.proto.jina_pb2_grpc import JinaDiscoverEndpointsRPCStub from jina.serve.executors import __dry_run_endpoint__ - from jina.serve.runtimes.helper import ( - _create_aux_model_doc_list_to_list, - _create_pydantic_model_from_schema, - ) + if not is_pydantic_v2: + from jina.serve.runtimes.helper import _create_aux_model_doc_list_to_list as create_pure_python_type_model + from jina.serve.runtimes.helper import _create_pydantic_model_from_schema as create_base_doc_from_schema + else: + from docarray.utils.create_dynamic_doc_class import create_pure_python_type_model, create_base_doc_from_schema channel = grpc.insecure_channel(f'0.0.0.0:{ports[0]}') stub = JinaDiscoverEndpointsRPCStub(channel) @@ -320,16 +322,16 @@ def bar(self, docs: DocList[Output1], **kwargs) -> DocList[Output2]: v = schema_map['/bar'] assert ( v['input'] - == _create_pydantic_model_from_schema( - _create_aux_model_doc_list_to_list(Input1).schema(), + == create_base_doc_from_schema( + create_pure_python_type_model(Input1).schema(), 'Input1', {}, ).schema() ) assert ( v['output'] - == _create_pydantic_model_from_schema( - _create_aux_model_doc_list_to_list(Output2).schema(), + == create_base_doc_from_schema( + create_pure_python_type_model(Output2).schema(), 'Output2', {}, ).schema() @@ -390,10 +392,11 @@ def bar(self, docs: DocList[Output1], **kwargs) -> DocList[Output2]: from jina.proto import jina_pb2 from jina.proto.jina_pb2_grpc import JinaDiscoverEndpointsRPCStub from jina.serve.executors import __default_endpoint__, __dry_run_endpoint__ - from jina.serve.runtimes.helper import ( - _create_aux_model_doc_list_to_list, - _create_pydantic_model_from_schema, - ) + if not is_pydantic_v2: + from jina.serve.runtimes.helper import _create_aux_model_doc_list_to_list as create_pure_python_type_model + from jina.serve.runtimes.helper import _create_pydantic_model_from_schema as create_base_doc_from_schema + else: + from docarray.utils.create_dynamic_doc_class import create_pure_python_type_model, create_base_doc_from_schema channel = grpc.insecure_channel(f'0.0.0.0:{ports[0]}') stub = JinaDiscoverEndpointsRPCStub(channel) @@ -411,14 +414,14 @@ def bar(self, docs: DocList[Output1], **kwargs) -> DocList[Output2]: v = schema_map[__default_endpoint__] assert ( v['input'] - == _create_pydantic_model_from_schema( - _create_aux_model_doc_list_to_list(Input1).schema(), 'Input1', {} + == create_base_doc_from_schema( + create_pure_python_type_model(Input1).schema(), 'Input1', {} ).schema() ) assert ( v['output'] - == _create_pydantic_model_from_schema( - _create_aux_model_doc_list_to_list(Output2).schema(), 'Output2', {} + == create_base_doc_from_schema( + create_pure_python_type_model(Output2).schema(), 'Output2', {} ).schema() ) diff --git a/tests/unit/serve/runtimes/test_helper.py b/tests/unit/serve/runtimes/test_helper.py index 323d0cdbc8b02..ac5dec8c593cb 100644 --- a/tests/unit/serve/runtimes/test_helper.py +++ b/tests/unit/serve/runtimes/test_helper.py @@ -3,7 +3,7 @@ import pytest -from jina._docarray import docarray_v2 +from jina._docarray import docarray_v2, is_pydantic_v2 from jina.serve.helper import get_default_grpc_options from jina.serve.runtimes.helper import ( _get_name_from_replicas_name, @@ -96,10 +96,11 @@ def test_create_pydantic_model_from_schema(transformation): from docarray.documents import TextDoc from docarray.typing import AnyTensor, ImageUrl - from jina.serve.runtimes.helper import ( - _create_aux_model_doc_list_to_list, - _create_pydantic_model_from_schema, - ) + if not is_pydantic_v2: + from jina.serve.runtimes.helper import _create_aux_model_doc_list_to_list as create_pure_python_type_model + from jina.serve.runtimes.helper import _create_pydantic_model_from_schema as create_base_doc_from_schema + else: + from docarray.utils.create_dynamic_doc_class import create_pure_python_type_model, create_base_doc_from_schema class Nested2Doc(BaseDoc): value: str @@ -124,8 +125,8 @@ class CustomDoc(BaseDoc): nested: Nested1Doc classvar: ClassVar[str] = 'classvar' - CustomDocCopy = _create_aux_model_doc_list_to_list(CustomDoc) - new_custom_doc_model = _create_pydantic_model_from_schema( + CustomDocCopy = create_pure_python_type_model(CustomDoc) + new_custom_doc_model = create_base_doc_from_schema( CustomDocCopy.schema(), 'CustomDoc', {} ) @@ -199,8 +200,8 @@ class CustomDoc(BaseDoc): class TextDocWithId(BaseDoc): ia: str - TextDocWithIdCopy = _create_aux_model_doc_list_to_list(TextDocWithId) - new_textdoc_with_id_model = _create_pydantic_model_from_schema( + TextDocWithIdCopy = create_pure_python_type_model(TextDocWithId) + new_textdoc_with_id_model = create_base_doc_from_schema( TextDocWithIdCopy.schema(), 'TextDocWithId', {} ) @@ -229,8 +230,8 @@ class TextDocWithId(BaseDoc): class ResultTestDoc(BaseDoc): matches: DocList[TextDocWithId] - ResultTestDocCopy = _create_aux_model_doc_list_to_list(ResultTestDoc) - new_result_test_doc_with_id_model = _create_pydantic_model_from_schema( + ResultTestDocCopy = create_pure_python_type_model(ResultTestDoc) + new_result_test_doc_with_id_model = create_base_doc_from_schema( ResultTestDocCopy.schema(), 'ResultTestDoc', {} ) result_test_docs = DocList[ResultTestDoc]( @@ -268,10 +269,11 @@ def test_create_empty_doc_list_from_schema(transformation): from docarray.documents import TextDoc from docarray.typing import AnyTensor, ImageUrl - from jina.serve.runtimes.helper import ( - _create_aux_model_doc_list_to_list, - _create_pydantic_model_from_schema, - ) + if not is_pydantic_v2: + from jina.serve.runtimes.helper import _create_aux_model_doc_list_to_list as create_pure_python_type_model + from jina.serve.runtimes.helper import _create_pydantic_model_from_schema as create_base_doc_from_schema + else: + from docarray.utils.create_dynamic_doc_class import create_pure_python_type_model, create_base_doc_from_schema class CustomDoc(BaseDoc): tensor: Optional[AnyTensor] @@ -288,8 +290,8 @@ class CustomDoc(BaseDoc): tags: Optional[Dict[str, Any]] = None lf: List[float] = [3.0, 4.1] - CustomDocCopy = _create_aux_model_doc_list_to_list(CustomDoc) - new_custom_doc_model = _create_pydantic_model_from_schema( + CustomDocCopy = create_pure_python_type_model(CustomDoc) + new_custom_doc_model = create_base_doc_from_schema( CustomDocCopy.schema(), 'CustomDoc', {} ) @@ -313,8 +315,8 @@ class CustomDoc(BaseDoc): class TextDocWithId(BaseDoc): ia: str - TextDocWithIdCopy = _create_aux_model_doc_list_to_list(TextDocWithId) - new_textdoc_with_id_model = _create_pydantic_model_from_schema( + TextDocWithIdCopy = create_pure_python_type_model(TextDocWithId) + new_textdoc_with_id_model = create_base_doc_from_schema( TextDocWithIdCopy.schema(), 'TextDocWithId', {} ) @@ -336,8 +338,8 @@ class TextDocWithId(BaseDoc): class ResultTestDoc(BaseDoc): matches: DocList[TextDocWithId] - ResultTestDocCopy = _create_aux_model_doc_list_to_list(ResultTestDoc) - new_result_test_doc_with_id_model = _create_pydantic_model_from_schema( + ResultTestDocCopy = create_pure_python_type_model(ResultTestDoc) + new_result_test_doc_with_id_model = create_base_doc_from_schema( ResultTestDocCopy.schema(), 'ResultTestDoc', {} ) result_test_docs = DocList[ResultTestDoc]() @@ -360,8 +362,11 @@ class ResultTestDoc(BaseDoc): @pytest.mark.skipif(not docarray_v2, reason='Test only working with docarray v2') def test_dynamic_class_creation_multiple_doclist_nested(): from docarray import BaseDoc, DocList - from jina.serve.runtimes.helper import _create_aux_model_doc_list_to_list - from jina.serve.runtimes.helper import _create_pydantic_model_from_schema + if not is_pydantic_v2: + from jina.serve.runtimes.helper import _create_aux_model_doc_list_to_list as create_pure_python_type_model + from jina.serve.runtimes.helper import _create_pydantic_model_from_schema as create_base_doc_from_schema + else: + from docarray.utils.create_dynamic_doc_class import create_pure_python_type_model, create_base_doc_from_schema class MyTextDoc(BaseDoc): text: str @@ -374,8 +379,8 @@ class SearchResult(BaseDoc): textlist = DocList[MyTextDoc]([MyTextDoc(text='hey')]) models_created_by_name = {} - SearchResult_aux = _create_aux_model_doc_list_to_list(SearchResult) - _ = _create_pydantic_model_from_schema( + SearchResult_aux = create_pure_python_type_model(SearchResult) + _ = create_base_doc_from_schema( SearchResult_aux.schema(), 'SearchResult', models_created_by_name ) QuoteFile_reconstructed_in_gateway_from_Search_results = models_created_by_name[