feat: fix batch handling in sagemaker for clip (#6216)

Co-authored-by: Jina Dev Bot <[email protected]>
jina-ai · Nov 4, 2024 · 92d03f0 · 92d03f0
1 parent 2f65ce1
commit 92d03f0
Show file tree

Hide file tree

Showing 5 changed files with 30 additions and 29 deletions.
diff --git a/jina/orchestrate/flow/base.py b/jina/orchestrate/flow/base.py
@@ -1776,10 +1776,8 @@ def build(self, copy_flow: bool = False, **kwargs) -> 'Flow':
             op_flow._deployment_nodes[GATEWAY_NAME].args.graph_description = json.dumps(
                 op_flow._get_graph_representation()
             )
-            op_flow._deployment_nodes[
-                GATEWAY_NAME
-            ].args.deployments_addresses = json.dumps(
-                op_flow._get_deployments_addresses()
+            op_flow._deployment_nodes[GATEWAY_NAME].args.deployments_addresses = (
+                json.dumps(op_flow._get_deployments_addresses())
             )
 
             op_flow._deployment_nodes[GATEWAY_NAME].update_pod_args()

diff --git a/jina/serve/executors/__init__.py b/jina/serve/executors/__init__.py
@@ -401,21 +401,21 @@ def __init__(
         self._init_monitoring()
         self._init_workspace = workspace
         if __dry_run_endpoint__ not in self.requests:
-            self.requests[
-                __dry_run_endpoint__
-            ] = _FunctionWithSchema.get_function_with_schema(
-                self.__class__._dry_run_func
+            self.requests[__dry_run_endpoint__] = (
+                _FunctionWithSchema.get_function_with_schema(
+                    self.__class__._dry_run_func
+                )
             )
         else:
             self.logger.warning(
                 f' Endpoint {__dry_run_endpoint__} is defined by the Executor. Be aware that this endpoint is usually reserved to enable health checks from the Client through the gateway.'
                 f' So it is recommended not to expose this endpoint. '
             )
         if type(self) == BaseExecutor:
-            self.requests[
-                __default_endpoint__
-            ] = _FunctionWithSchema.get_function_with_schema(
-                self.__class__._dry_run_func
+            self.requests[__default_endpoint__] = (
+                _FunctionWithSchema.get_function_with_schema(
+                    self.__class__._dry_run_func
+                )
             )
 
         self._lock = contextlib.AsyncExitStack()
@@ -595,14 +595,14 @@ def _add_requests(self, _requests: Optional[Dict]):
                 _func = getattr(self.__class__, func)
                 if callable(_func):
                     # the target function is not decorated with `@requests` yet
-                    self.requests[
-                        endpoint
-                    ] = _FunctionWithSchema.get_function_with_schema(_func)
+                    self.requests[endpoint] = (
+                        _FunctionWithSchema.get_function_with_schema(_func)
+                    )
                 elif typename(_func) == 'jina.executors.decorators.FunctionMapper':
                     # the target function is already decorated with `@requests`, need unwrap with `.fn`
-                    self.requests[
-                        endpoint
-                    ] = _FunctionWithSchema.get_function_with_schema(_func.fn)
+                    self.requests[endpoint] = (
+                        _FunctionWithSchema.get_function_with_schema(_func.fn)
+                    )
                 else:
                     raise TypeError(
                         f'expect {typename(self)}.{func} to be a function, but receiving {typename(_func)}'

diff --git a/jina/serve/runtimes/worker/http_csp_app.py b/jina/serve/runtimes/worker/http_csp_app.py
@@ -182,13 +182,14 @@ def construct_model_from_line(
                         # Handle list of nested models
                         elif get_origin(field_type) is list:
                             list_item_type = get_args(field_type)[0]
-                            parsed_list = json.loads(field_str)
-                            if issubclass(list_item_type, BaseModel):
-                                parsed_fields[field_name] = parse_obj_as(
-                                    List[list_item_type], parsed_list
-                                )
-                            else:
-                                parsed_fields[field_name] = parsed_list
+                            if field_str:
+                                parsed_list = json.loads(field_str)
+                                if issubclass(list_item_type, BaseModel):
+                                    parsed_fields[field_name] = parse_obj_as(
+                                        List[list_item_type], parsed_list
+                                    )
+                                else:
+                                    parsed_fields[field_name] = parsed_list
                         # General parsing attempt for other types
                         else:
                             if field_str:

diff --git a/tests/integration/docarray_v2/csp/SampleClipExecutor/executor.py b/tests/integration/docarray_v2/csp/SampleClipExecutor/executor.py
@@ -1,4 +1,4 @@
-from typing import Optional
+from typing import Optional, List
 
 import numpy as np
 from docarray import BaseDoc, DocList
@@ -13,6 +13,8 @@ class TextAndImageDoc(BaseDoc):
     text: Optional[str] = None
     url: Optional[AnyUrl] = None
     bytes: Optional[ImageBytes] = None
+    num_tokens: Optional[int] = None
+    input_ids: Optional[List[int]] = None
 
 
 class EmbeddingResponseModel(TextAndImageDoc):

diff --git a/tests/integration/docarray_v2/csp/valid_clip_input.csv b/tests/integration/docarray_v2/csp/valid_clip_input.csv