feat: support param in sm batch (#6229)

Co-authored-by: Joan Martinez <[email protected]>

Author: zac-li

.github/workflows/build-old-docs.yml

@@ -92,7 +92,7 @@ jobs:
           mv ./docs/_build/dirhtml ./${{ matrix.version }}
           zip -r /tmp/build.zip ./${{ matrix.version }}/*
       - name: Upload built html
-        uses: actions/upload-artifact@v3
+        uses: actions/upload-artifact@v4
         with:
             name: ${{ matrix.version }}
             path: /tmp/build.zip
@@ -106,7 +106,7 @@ jobs:
         with:
           fetch-depth: 1
           ref: ${{ inputs.pages_branch }}
-      - uses: actions/download-artifact@v3
+      - uses: actions/download-artifact@v4
         with:
           path: /tmp/artifacts
       - name: Clear old builds

.github/workflows/cd.yml

@@ -245,7 +245,7 @@ jobs:
     runs-on: ubuntu-latest
     steps:
       - name: Check out repository
-        uses: actions/checkout@v2
+        uses: actions/checkout@v2.5.0
         with:
           fetch-depth: 200
 
@@ -347,7 +347,7 @@ jobs:
             goarch: arm64
     steps:
       - name: Check out repository
-        uses: actions/checkout@v2
+        uses: actions/checkout@v2.5.0
 
       - name: Set up Python ${{ matrix.python }}
         uses: actions/setup-python@v2
@@ -399,8 +399,9 @@ jobs:
           python -m cibuildwheel --output-dir dist
 
       - name: Upload wheels as artifacts
-        uses: actions/upload-artifact@v3
+        uses: actions/upload-artifact@v4
         with:
+          name: artifacts-${{ strategy.job-index }}
           path: dist/*.whl
 
 # comment for now, do it manually if needed

.github/workflows/ci.yml

@@ -651,7 +651,7 @@ jobs:
 
     steps:
       - name: Check out repository
-        uses: actions/checkout@v2
+        uses: actions/checkout@v2.5.0
 
       - name: Set up Python ${{ matrix.python }}
         uses: actions/setup-python@v2
@@ -688,8 +688,9 @@ jobs:
           python -m cibuildwheel --output-dir dist
 
       - name: Upload wheels as artifacts
-        uses: actions/upload-artifact@v3
+        uses: actions/upload-artifact@v4
         with:
+          name: artifacts-${{ strategy.job-index }}
           path: dist/*.whl
 
   core-test:

.github/workflows/force-release.yml

@@ -86,7 +86,7 @@ jobs:
 
     steps:
       - name: Check out repository
-        uses: actions/checkout@v2
+        uses: actions/checkout@v2.5.0
         with:
           fetch-depth: 200
 
@@ -124,8 +124,9 @@ jobs:
           python -m cibuildwheel --output-dir dist
           
       - name: Upload wheels as artifacts
-        uses: actions/upload-artifact@v3
+        uses: actions/upload-artifact@v4
         with:
+          name: artifacts-${{ strategy.job-index }}
           path: dist/*.whl
 
   regular-release:
@@ -141,9 +142,10 @@ jobs:
         with:
           python-version: "3.10"
         # https://github.com/actions/checkout#fetch-all-tags
-      - uses: actions/download-artifact@v3
+      - uses: actions/download-artifact@v4
         with:
-          name: artifact
+          pattern: artifacts-*
+          merge-multiple: true
           path: dist
       - run: |
           git fetch --depth=200

jina/serve/runtimes/worker/http_csp_app.py

@@ -1,4 +1,5 @@
-from typing import TYPE_CHECKING, Callable, Dict, List, Optional, Union
+from typing import (TYPE_CHECKING, Callable, Dict, List, Literal, Optional,
+                    Union)
 
 from jina._docarray import docarray_v2
 from jina.importer import ImportExtensions
@@ -74,7 +75,7 @@ def add_post_route(
             input_model,
             output_model,
             input_doc_list_model=None,
-            output_doc_list_model=None,
+            parameters_model=None,
     ):
         import json
         from typing import List, Type, Union
@@ -150,54 +151,85 @@ async def post(request: Request):
                 csv_body = bytes_body.decode('utf-8')
                 if not is_valid_csv(csv_body):
                     raise HTTPException(
-                        status_code=400,
+                        status_code=http_status.HTTP_400_BAD_REQUEST,
                         detail='Invalid CSV input. Please check your input.',
                     )
 
-                def construct_model_from_line(
-                        model: Type[BaseModel], line: List[str]
-                ) -> BaseModel:
+
+                def construct_model_from_line(model: Type[BaseModel], line: List[str]) -> BaseModel:
+                    origin = get_origin(model)
+                    # If the model is of type Optional[X], unwrap it to get X
+                    if origin is Union:
+                        # If the model is of type Optional[X], unwrap it to get X
+                        args = get_args(model)
+                        if type(None) in args:
+                            model = args[0]
+
                     parsed_fields = {}
                     model_fields = model.__fields__
 
-                    for field_str, (field_name, field_info) in zip(
-                            line, model_fields.items()
-                    ):
-                        field_type = field_info.outer_type_
-
-                        # Handle Union types by attempting to arse each potential type
-                        if get_origin(field_type) is Union:
-                            for possible_type in get_args(field_type):
-                                if possible_type is str:
-                                    parsed_fields[field_name] = field_str
-                                    break
-                                else:
+                    for idx, (field_name, field_info) in enumerate(model_fields.items()):
+                        field_type = field_info.type_
+                        field_str = line[idx]  # Corresponding value from the row
+
+                        try:
+                            # Handle Literal types (e.g., Optional[Literal["value1", "value2"]])
+                            origin = get_origin(field_type)
+                            if origin is Literal:
+                                literal_values = get_args(field_type)
+                                if field_str not in literal_values:
+                                    raise HTTPException(
+                                        status_code=http_status.HTTP_400_BAD_REQUEST,
+                                        detail=f"Invalid value '{field_str}' for field '{field_name}'. Expected one of: {literal_values}"
+                                    )
+                                parsed_fields[field_name] = field_str
+
+                            # Handle Union types (e.g., Optional[int, str])
+                            elif origin is Union:
+                                for possible_type in get_args(field_type):
                                     try:
-                                        parsed_fields[field_name] = parse_obj_as(
-                                            possible_type, json.loads(field_str)
-                                        )
+                                        parsed_fields[field_name] = parse_obj_as(possible_type, field_str)
                                         break
-                                    except (json.JSONDecodeError, ValidationError):
+                                    except (ValueError, TypeError, ValidationError):
                                         continue
-                        # Handle list of nested models
-                        elif get_origin(field_type) is list:
-                            list_item_type = get_args(field_type)[0]
-                            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:
-                                try:
-                                    parsed_fields[field_name] = field_info.type_(field_str)
-                                except (ValueError, TypeError):
-                                    # Fallback to parse_obj_as when type is more complex, e., AnyUrl or ImageBytes
-                                    parsed_fields[field_name] = parse_obj_as(field_info.type_, field_str)
+
+                            # Handle list of nested models (e.g., List[Item])
+                            elif get_origin(field_type) is list:
+                                list_item_type = get_args(field_type)[0]
+                                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
+
+                            # Handle other general types
+                            else:
+                                if field_str:
+                                    if field_type == bool:
+                                        # Special case: handle "false" and "true" as booleans
+                                        if field_str.lower() == "false":
+                                            parsed_fields[field_name] = False
+                                        elif field_str.lower() == "true":
+                                            parsed_fields[field_name] = True
+                                        else:
+                                            raise HTTPException(
+                                                status_code=http_status.HTTP_400_BAD_REQUEST,
+                                                detail=f"Invalid value '{field_str}' for boolean field '{field_name}'. Expected 'true' or 'false'."
+                                            )
+                                    else:
+                                        # General case: try converting to the target type
+                                        try:
+                                            parsed_fields[field_name] = field_type(field_str)
+                                        except (ValueError, TypeError):
+                                            # Fallback to parse_obj_as when type is more complex, e., AnyUrl or ImageBytes
+                                            parsed_fields[field_name] = parse_obj_as(field_type, field_str)
+
+                        except Exception as e:
+                            raise HTTPException(
+                                status_code=http_status.HTTP_400_BAD_REQUEST,
+                                detail=f"Error parsing value '{field_str}' for field '{field_name}': {str(e)}"
+                            )
 
                     return model(**parsed_fields)
 
@@ -209,25 +241,41 @@ def construct_model_from_line(
                 # We also expect the csv file to have no quotes and use the escape char '\'
                 field_names = [f for f in input_doc_list_model.__fields__]
                 data = []
+                parameters = None
+                first_row = True
                 for line in csv.reader(
                         StringIO(csv_body),
                         delimiter=',',
                         quoting=csv.QUOTE_NONE,
                         escapechar='\\',
                 ):
-                    if len(line) != len(field_names):
-                        raise HTTPException(
-                            status_code=400,
-                            detail=f'Invalid CSV format. Line {line} doesn\'t match '
-                                   f'the expected field order {field_names}.',
-                        )
-                    data.append(construct_model_from_line(input_doc_list_model, line))
-
-                return await process(input_model(data=data))
+                    if first_row:
+                        first_row = False
+                        if len(line) > 1 and line[1] == 'params_row':  # Check if it's a parameters row by examining the 2nd text in the first line
+                            parameters = construct_model_from_line(parameters_model, line[2:])
+                        else:
+                            if len(line) != len(field_names):
+                                raise HTTPException(
+                                    status_code=http_status.HTTP_400_BAD_REQUEST,
+                                    detail=f'Invalid CSV format. Line {line} doesn\'t match '
+                                        f'the expected field order {field_names}.',
+                                )
+                            data.append(construct_model_from_line(input_doc_list_model, line))
+                    else:
+                        # Treat it as normal data row
+                        if len(line) != len(field_names):
+                            raise HTTPException(
+                                status_code=http_status.HTTP_400_BAD_REQUEST,
+                                detail=f'Invalid CSV format. Line {line} doesn\'t match '
+                                    f'the expected field order {field_names}.',
+                            )
+                        data.append(construct_model_from_line(input_doc_list_model, line))
+
+                return await process(input_model(data=data, parameters=parameters))
 
             else:
                 raise HTTPException(
-                    status_code=400,
+                    status_code=http_status.HTTP_400_BAD_REQUEST,
                     detail=f'Invalid content-type: {content_type}. '
                            f'Please use either application/json or text/csv.',
                 )
@@ -273,7 +321,7 @@ def construct_model_from_line(
                 input_model=endpoint_input_model,
                 output_model=endpoint_output_model,
                 input_doc_list_model=input_doc_model,
-                output_doc_list_model=output_doc_model,
+                parameters_model=parameters_model,
             )
 
     from jina.serve.runtimes.gateway.health_model import JinaHealthModel

tests/integration/docarray_v2/csp/SampleExecutor/executor.py

@@ -2,6 +2,7 @@
 from docarray import BaseDoc, DocList
 from docarray.typing import NdArray
 from pydantic import Field, BaseModel
+from typing import Optional, Literal
 
 from jina import Executor, requests
 
@@ -20,7 +21,17 @@ class Config(BaseDoc.Config):
 
 
 class Parameters(BaseModel):
-    emb_dim: int
+    task: Optional[
+        Literal[
+            "retrieval.query",
+            "retrieval.passage",
+            "text-matching",
+            "classification",
+            "separation",
+        ]
+    ] = None
+    late_chunking: bool = False
+    dimensions: Optional[int] = None
 
 
 
@@ -46,7 +57,7 @@ def bar(self, docs: DocList[TextDoc], parameters: Parameters, **kwargs) -> DocLi
                 EmbeddingResponseModel(
                     id=doc.id,
                     text=doc.text,
-                    embeddings=np.random.random((1, parameters.emb_dim)),
+                    embeddings=np.random.random((1, parameters.dimensions)),
                 )
             )
         return DocList[EmbeddingResponseModel](ret)

tests/integration/docarray_v2/csp/test_sagemaker_embedding.py

@@ -89,7 +89,7 @@ def test_provider_sagemaker_pod_inference_parameters():
                     'data': [
                         {'text': 'hello world'},
                     ],
-                    'parameters': {'emb_dim': emb_dim}
+                    'parameters': {'dimensions': emb_dim}
                 },
             )
             assert resp.status_code == 200
@@ -148,6 +148,50 @@ def test_provider_sagemaker_pod_batch_transform_valid(filename):
             assert len(d["embeddings"][0]) == 64
 
 
+def test_provider_sagemaker_pod_batch_transform_with_params_valid():
+    args, _ = set_pod_parser().parse_known_args(
+        [
+            '--uses',
+            os.path.join(os.path.dirname(__file__), "SampleExecutor", "config.yml"),
+            '--provider',
+            'sagemaker',
+            "--provider-endpoint",
+            "encode_parameter",
+            'serve',  # This is added by sagemaker
+        ]
+    )
+    with Pod(args):
+        texts = []
+        with open(os.path.join(os.path.dirname(__file__), "valid_input_3.csv"), "r") as f:
+            csv_data = f.read()
+
+        csv_reader = csv.reader(io.StringIO(csv_data), delimiter=",", quoting=csv.QUOTE_NONE, escapechar="\\")
+
+        # Before comparison, remove the parameters row
+        next(csv_reader)
+
+        for line in csv_reader:
+            texts.append(line[1])
+
+        resp = requests.post(
+            f"http://localhost:{sagemaker_port}/invocations",
+            headers={
+                "accept": "application/json",
+                "content-type": "text/csv",
+            },
+            data=csv_data,
+        )
+        assert resp.status_code == 200
+        resp_json = resp.json()
+        assert len(resp_json["data"]) == 10
+        for idx, d in enumerate(resp_json["data"]):
+            assert d["text"] == texts[idx]
+            assert len(d["embeddings"][0]) == 2
+
+        assert resp_json["parameters"]["late_chunking"] == False
+        assert resp_json["parameters"]["task"] == "retrieval.query"
+
+
 def test_provider_sagemaker_pod_batch_transform_invalid():
     args, _ = set_pod_parser().parse_known_args(
         [