Repository: mosecorg/mosec (language: python)

Issue:
feat: type check/validation for inference parameters
### Describe the feature

Users will get confused with "500, 'inference internal error" when they send a wrong format parameters.
We could use some type hint check to announce users the right format.

## An idea from `pydantic`
```python
from functools import wraps
from mosec import ValidationError, Worker
from mosec.mixin import MsgpackMixin
from typing import Any, Dict, List, Literal, TypedDict
from pydantic import create_model_from_typeddict
import pydantic

# exported from mosec
class data_type_check(object):
    def __init__(self, data_type: type):
        self.data_type = data_type

    def __call__(self, func):
        @wraps(func)
        def wrapped_function(cls, data: Any):
            if not isinstance(data, self.data_type):
                # could extend ValidationError and display self.data_type
                raise ValidationError
            if isinstance(data, dict):
                self.validate_dict(data)
            # skip list validation for arbitrary length
            return func(cls, data)
        return wrapped_function

    def validate_dict(self, data: Dict):
        Model = create_model_from_typeddict(self.data_type)
        try:
            Model.parse_obj(data)
        except pydantic.error_wrappers.ValidationError:
            # could extend ValidationError and display Model.schema_json()
            raise ValidationError

# user code
class Input(TypedDict):
    task: Literal["image", "audio", "video", "text"]
    data: List[bytes | str]

class ImageBind(MsgpackMixin, Worker):
    def __init__(self):
        pass

    @data_type_check(data_type=Input)
    def forward(self, data: Input) -> List[List[float]]:
        pass
```



### Why do you need this feature?

Terminal user will get benefit from type validation, if model developer take a little time to define a schema.

### Additional context

_No response_


Discussion / hints:
You don't need to write a decorator. Just use the mixin to override the `deserialize` method. It's more efficient to parse the raw bytes into the desired type than parsing an object after the deserialization.

BTW, in this case, [msgspec](https://jcristharif.com/msgspec/benchmarks.html) might be faster than pydantic v1.
I hope this feature is optional. We do not run type checks if the user does not add any type hint in the server code.
Actually, typed input can also generate the OpenAPI doc. It's supported by pydantic/msgspec/etc.
> Actually, typed input can also generate the OpenAPI doc. It's supported by pydantic/msgspec/etc.

It's also useful to users. I think. We could generate the API doc from the server code.
> You don't need to write a decorator. Just use the mixin to override the `deserialize` method. It's more efficient to parse the raw bytes into the desired type than parsing an object after the deserialization.
> 
> BTW, in this case, [msgspec](https://jcristharif.com/msgspec/benchmarks.html) might be faster than pydantic v1.

It's designed for endpoint users, when they write a request of the wrong format, they could receive some suggestion of the right `type schema` in response, instead of `inference internal error`.

`pydantic` could provide some schema json from its model, likes:
```json
{
	"title": "Input",
	"type": "object",
	"properties": {
		"task": {
			"title": "Task",
			"enum": ["image", "audio", "video", "text"],
			"type": "string"
		},
		"data": {
			"title": "Data",
			"type": "array",
			"items": {
				"anyOf": [{
					"type": "string",
					"format": "binary"
				}, {
					"type": "string"
				}]
			}
		}
	},
	"required": ["task", "data"]
}
```

This case usually happens when client and server are coded by different people.
I don't see any conflicts between your comments and the quoted comments.
> I don't see any conflicts between your comments and the quoted comments.

Sorry, I got it recently by read the doc. `msgspec` could do same thing and really faster.
Maybe we could insert an msgspec-based validator and response a revised `ValidateError` in mosec when validating failed? 
It is optional by server developer.

> Maybe we could insert an msgspec-based validator and response a revised `ValidateError` in mosec when validating failed? It is optional by server developer.

Yes, we can do that. There is a `ValidationError` for this purpose because mosec used to have a pydantic validator (not in this repo history). I think it can be one of the mixin. You can try it. Feel free to ask questions.
The mixin might be a little complex due to the multiprocessing spawn context. Let me create a draft for you. But feel free to review and leave some comments.


Fix patch:
diff --git a/.github/workflows/check.yml b/.github/workflows/check.yml
index ac837fe2..8b21cd31 100644
--- a/.github/workflows/check.yml
+++ b/.github/workflows/check.yml
@@ -59,7 +59,7 @@ jobs:
     strategy:
       max-parallel: 10
       matrix:
-        python-version: [3.7, 3.8, 3.9, "3.10", "3.11"]
+        python-version: [3.8, 3.9, "3.10", "3.11"]
         os: [ubuntu-latest , macos-latest]
 
     steps:
@@ -74,7 +74,7 @@ jobs:
       - name: Install dependencies
         run: |
           python -m pip install --upgrade pip
-          pip install -e .[dev,plugin]
+          pip install -e .[dev,mixin]
           make dev
       - name: Test
         run: |
diff --git a/.github/workflows/nightly.yml b/.github/workflows/nightly.yml
index be9583b8..910ddd2a 100644
--- a/.github/workflows/nightly.yml
+++ b/.github/workflows/nightly.yml
@@ -30,7 +30,7 @@ jobs:
       - name: Install dependencies
         run: |
           python -m pip install --upgrade pip
-          pip install -e .[dev,plugin]
+          pip install -e .[dev,mixin]
           make dev
       - name: Test
         run: |
diff --git a/.github/workflows/page.yml b/.github/workflows/page.yml
index c1ccd0f3..2b3b0761 100644
--- a/.github/workflows/page.yml
+++ b/.github/workflows/page.yml
@@ -42,7 +42,7 @@ jobs:
     - name: Install dependencies
       run: |
         python -m pip install --upgrade pip
-        pip install -e .[dev,doc,plugin]
+        pip install -e .[dev,doc,mixin]
     - name: Generate docs
       run: |
         cd docs && make html
diff --git a/.gitignore b/.gitignore
index aefd8d82..eaebf965 100644
--- a/.gitignore
+++ b/.gitignore
@@ -144,3 +144,6 @@ target/
 
 # version file generated by setuptools_scm
 mosec/_version.py
+
+# ruff
+.ruff_cache/
diff --git a/Makefile b/Makefile
index 7969bb97..71efefb3 100644
--- a/Makefile
+++ b/Makefile
@@ -4,7 +4,7 @@ PY_SOURCE_FILES=mosec tests examples setup.py
 RUST_SOURCE_FILES=src/*
 
 install:
-	pip install -e .[dev,doc,plugin]
+	pip install -e .[dev,doc,mixin]
 	pre-commit install
 	rustup toolchain install nightly
 	rustup component add rustfmt --toolchain nightly
@@ -21,13 +21,13 @@ test: dev
 	RUST_BACKTRACE=1 cargo test -vv
 
 test_arrow: dev
-	@pip install -q -r requirements/plugin.txt
-	echo "Running tests for the plugin"
+	@pip install -q -r requirements/mixin.txt
+	echo "Running tests for the mixin"
 	pytest tests -vv -s -m "arrow"
-	pip uninstall -y -r requirements/plugin.txt
+	pip uninstall -y -r requirements/mixin.txt
 
 test_all: dev
-	@pip install -q -r requirements/plugin.txt
+	@pip install -q -r requirements/mixin.txt
 	echo "Running tests for the all features"
 	pytest tests -vv -s
 	RUST_BACKTRACE=1 cargo test -vv
diff --git a/README.md b/README.md
index 3ee53bf9..d4fc231e 100644
--- a/README.md
+++ b/README.md
@@ -163,6 +163,7 @@ That's it! You have just hosted your **_stable-diffusion model_** as a service!
 More ready-to-use examples can be found in the [Example](https://mosecorg.github.io/mosec/examples/index.html) section. It includes:
 
 - [Multi-stage workflow demo](https://mosecorg.github.io/mosec/examples/echo.html): a simple echo demo even without any ML model.
+- [Request validation](https://mosecorg.github.io/mosec/examples/validate.html): validate the request with type annotation.
 - [Shared memory IPC](https://mosecorg.github.io/mosec/examples/ipc.html): inter-process communication with shared memory.
 - [Customized GPU allocation](https://mosecorg.github.io/mosec/examples/env.html): deploy multiple replicas, each using different GPUs.
 - [Customized metrics](https://mosecorg.github.io/mosec/examples/metric.html): record your own metrics for monitoring.
diff --git a/docs/source/examples/index.md b/docs/source/examples/index.md
index 9f8e25f3..828a6814 100644
--- a/docs/source/examples/index.md
+++ b/docs/source/examples/index.md
@@ -12,6 +12,7 @@ jax
 metric
 pytorch
 stable_diffusion
+validate
 ```
 
 We provide examples across different ML frameworks and for various tasks in this section.
diff --git a/docs/source/examples/validate.md b/docs/source/examples/validate.md
new file mode 100644
index 00000000..0ce55519
--- /dev/null
+++ b/docs/source/examples/validate.md
@@ -0,0 +1,49 @@
+# Validate Request
+
+This example shows how to use the `TypedMsgPackMixin` to validate the request with the help of [`msgspec`](https://github.com/jcrist/msgspec).
+
+Request validation can provide the following benefits:
+
+- The client can know the exact expected data schema from the type definition.
+- Validation failure will return the details of the failure reason to help the client debug.
+- Ensure that the service is working on the correct data without fear.
+
+First of all, define the request type with `msgspec.Struct` like:
+
+```python
+class Request(msgspec.Struct):
+    media: str
+    binary: bytes
+```
+
+Then, apply the `TypedMsgPackMixin` mixin and add the type you defined to the annotation of `forward(self, data)`:
+
+```python
+class Inference(TypedMsgPackMixin, Worker):
+    def forward(self, data: Request):
+        pass
+```
+
+```{note}
+If you are using dynamic **batch** inference as the first stage, just use the `List[Request]` as the annotation.
+```
+
+You can check the full demo code below.
+
+## Server
+
+```{include} ../../../examples/type_validation/server.py
+:code: python
+```
+
+## Client
+
+```{include} ../../../examples/type_validation/client.py
+:code: python
+```
+
+## Test
+
+```shell
+python client.py
+```
diff --git a/examples/resnet50_client_msgpack.py b/examples/resnet50_client_msgpack.py
index ea119492..004c6531 100644
--- a/examples/resnet50_client_msgpack.py
+++ b/examples/resnet50_client_msgpack.py
@@ -23,7 +23,7 @@
 
 prediction = httpx.post(
     "http://127.0.0.1:8000/inference",
-    data=msgpack.packb({"image": dog_bytes}),
+    content=msgpack.packb({"image": dog_bytes}),
 )
 if prediction.status_code == 200:
     print(msgpack.unpackb(prediction.content))
diff --git a/examples/stable_diffusion/client.py b/examples/stable_diffusion/client.py
index 398ef0cc..c75147d0 100644
--- a/examples/stable_diffusion/client.py
+++ b/examples/stable_diffusion/client.py
@@ -37,7 +37,7 @@
 args = parser.parse_args()
 resp = httpx.post(
     f"http://127.0.0.1:{args.port}/inference",
-    data=msgpack.packb(args.prompt),
+    content=msgpack.packb(args.prompt),
     timeout=httpx.Timeout(20),
 )
 if resp.status_code == 200:
diff --git a/examples/type_validation/client.py b/examples/type_validation/client.py
new file mode 100644
index 00000000..9be4261e
--- /dev/null
+++ b/examples/type_validation/client.py
@@ -0,0 +1,29 @@
+# Copyright 2023 MOSEC Authors
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from http import HTTPStatus
+
+import httpx
+import msgpack  # type: ignore
+
+req = {
+    "bin": b"hello mosec",
+    "name": "type check",
+}
+
+resp = httpx.post("http://127.0.0.1:8000/inference", content=msgpack.packb(req))
+if resp.status_code == HTTPStatus.OK:
+    print(f"OK: {msgpack.unpackb(resp.content)}")
+else:
+    print(f"err[{resp.status_code}] {resp.text}")
diff --git a/examples/type_validation/server.py b/examples/type_validation/server.py
new file mode 100644
index 00000000..e60fc459
--- /dev/null
+++ b/examples/type_validation/server.py
@@ -0,0 +1,54 @@
+# Copyright 2023 MOSEC Authors
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""Request validation example."""
+
+from typing import Any, List
+
+from msgspec import Struct
+
+from mosec import Server, Worker
+from mosec.mixin import TypedMsgPackMixin
+
+
+class Request(Struct):
+    """User request struct."""
+
+    # pylint: disable=too-few-public-methods
+
+    bin: bytes
+    name: str = "test"
+
+
+class Preprocess(TypedMsgPackMixin, Worker):
+    """Dummy preprocess to exit early if the validation failed."""
+
+    def forward(self, data: Request) -> Any:
+        """Input will be parse as the `Request`."""
+        print(f"received {data}")
+        return data.bin
+
+
+class Inference(TypedMsgPackMixin, Worker):
+    """Dummy batch inference."""
+
+    def forward(self, data: List[bytes]) -> Any:
+        return [len(buf) for buf in data]
+
+
+if __name__ == "__main__":
+    server = Server()
+    server.append_worker(Preprocess)
+    server.append_worker(Inference, max_batch_size=16)
+    server.run()
diff --git a/mosec/mixin/__init__.py b/mosec/mixin/__init__.py
index bb4a45ea..4479a99f 100644
--- a/mosec/mixin/__init__.py
+++ b/mosec/mixin/__init__.py
@@ -17,5 +17,6 @@
 from mosec.mixin.msgpack_worker import MsgpackMixin
 from mosec.mixin.numbin_worker import NumBinIPCMixin
 from mosec.mixin.plasma_worker import PlasmaShmIPCMixin
+from mosec.mixin.typed_worker import TypedMsgPackMixin
 
-__all__ = ["MsgpackMixin", "NumBinIPCMixin", "PlasmaShmIPCMixin"]
+__all__ = ["MsgpackMixin", "NumBinIPCMixin", "PlasmaShmIPCMixin", "TypedMsgPackMixin"]
diff --git a/mosec/mixin/typed_worker.py b/mosec/mixin/typed_worker.py
new file mode 100644
index 00000000..f2eaaeef
--- /dev/null
+++ b/mosec/mixin/typed_worker.py
@@ -0,0 +1,79 @@
+# Copyright 2023 MOSEC Authors
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""MOSEC type validation mixin."""
+
+import inspect
+from typing import Any, List
+
+import msgspec
+
+from mosec.errors import ValidationError
+from mosec.worker import Worker
+
+
+def parse_forward_input_type(func):
+    """Parse the input type of the forward function.
+
+    - single request: return the type
+    - batch request: return the list item type
+    """
+    sig = inspect.signature(func)
+    params = list(sig.parameters.values())
+    if len(params) < 1:
+        raise TypeError("`forward` method doesn't have enough(1) parameters")
+
+    typ = params[0].annotation
+    origin = getattr(typ, "__origin__", None)
+    if origin is None:
+        return typ
+    # GenericAlias, `func` could be batch inference
+    if origin is list or origin is List:
+        if not hasattr(typ, "__args__") or len(typ.__args__) != 1:
+            raise TypeError(
+                "`forward` with dynamic batch should use "
+                "`List[Struct]` as the input annotation"
+            )
+        return typ.__args__[0]
+    raise TypeError(f"unsupported type {typ}")
+
+
+class TypedMsgPackMixin(Worker):
+    """Enable request type validation with `msgspec` and serde with `msgpack`."""
+
+    # pylint: disable=no-self-use
+
+    _input_type = None
+
+    def _get_input_type(self):
+        """Get the input type from annotations."""
+        if self._input_type is None:
+            self._input_type = parse_forward_input_type(self.forward)
+        return self._input_type
+
+    def deserialize(self, data: Any) -> bytes:
+        """Deserialize and validate request with msgspec."""
+        schema = self._get_input_type()
+        if not issubclass(schema, msgspec.Struct):
+            # skip other annotation type
+            return super().deserialize(data)
+
+        try:
+            return msgspec.msgpack.decode(data, type=schema)
+        except msgspec.ValidationError as err:
+            raise ValidationError(err)  # pylint: disable=raise-missing-from
+
+    def serialize(self, data: Any) -> bytes:
+        """Serialize with `msgpack`."""
+        return msgspec.msgpack.encode(data)
diff --git a/pyproject.toml b/pyproject.toml
index 345ba29e..7c3ef593 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -17,7 +17,6 @@ classifiers = [
     "Intended Audience :: Science/Research",
     "License :: OSI Approved :: Apache Software License",
     "Programming Language :: Python :: 3 :: Only",
-    "Programming Language :: Python :: 3.7",
     "Programming Language :: Python :: 3.8",
     "Programming Language :: Python :: 3.9",
     "Programming Language :: Python :: 3.10",
@@ -60,7 +59,7 @@ pretty = true
 
 [tool.pyright]
 pythonPlatform = "All"
-pythonVersion = "3.7"
+pythonVersion = "3.8"
 include = ["mosec", "tests", "examples"]
 
 [tool.isort]
diff --git a/requirements/dev.txt b/requirements/dev.txt
index 24ba9fdb..4a37bfcf 100644
--- a/requirements/dev.txt
+++ b/requirements/dev.txt
@@ -9,4 +9,5 @@ black>=20.8.0
 isort>=5.6
 autoflake>=1.4
 pre-commit>=2.15.0
+msgpack>=1.0.5
 httpx==0.24.0
diff --git a/requirements/plugin.txt b/requirements/mixin.txt
similarity index 75%
rename from requirements/plugin.txt
rename to requirements/mixin.txt
index 065b0981..ba817f7d 100644
--- a/requirements/plugin.txt
+++ b/requirements/mixin.txt
@@ -1,3 +1,4 @@
-pyarrow>=0.6.1,<12
 msgpack>=1.0.2
+msgspec>=0.15.0
 numbin>=0.4.0
+pyarrow>=0.6.1,<12
diff --git a/setup.py b/setup.py
index a948fec8..cd82c1f4 100644
--- a/setup.py
+++ b/setup.py
@@ -18,8 +18,8 @@
 with open(os.path.join(here, "requirements/doc.txt"), encoding="utf-8") as f:
     doc_requirements = f.read().splitlines()
 
-with open(os.path.join(here, "requirements/plugin.txt"), encoding="utf-8") as f:
-    plugin_requirements = f.read().splitlines()
+with open(os.path.join(here, "requirements/mixin.txt"), encoding="utf-8") as f:
+    mixin_requirements = f.read().splitlines()
 
 
 class RustExtension(Extension):
@@ -80,11 +80,11 @@ def build_extension(self, ext: Extension):
     use_scm_version=True,
     packages=find_packages(exclude=["examples*", "tests*"]),
     include_package_data=True,
-    python_requires=">=3.6",
+    python_requires=">=3.8",
     setup_requires=["setuptools_scm>=7.0"],
     extras_require={
         "dev": dev_requirements,
-        "plugin": plugin_requirements,
+        "mixin": mixin_requirements,
         "doc": doc_requirements,
     },
     zip_safe=False,


Test patch:
diff --git a/tests/mixin_typed_service.py b/tests/mixin_typed_service.py
new file mode 100644
index 00000000..8bbca4e0
--- /dev/null
+++ b/tests/mixin_typed_service.py
@@ -0,0 +1,36 @@
+# Copyright 2023 MOSEC Authors
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from typing import Any, List
+
+from msgspec import Struct
+
+from mosec import Server, Worker
+from mosec.mixin import TypedMsgPackMixin
+
+
+class Request(Struct):
+    media: str
+    binary: bytes
+
+
+class Inference(TypedMsgPackMixin, Worker):
+    def forward(self, data: List[Request]) -> Any:
+        return [len(req.binary) for req in data]
+
+
+if __name__ == "__main__":
+    server = Server()
+    server.append_worker(Inference, max_batch_size=4)
+    server.run()
diff --git a/tests/test_service.py b/tests/test_service.py
index 9c87ae95..ac368ee1 100644
--- a/tests/test_service.py
+++ b/tests/test_service.py
@@ -21,6 +21,7 @@
 from threading import Thread
 
 import httpx
+import msgpack  # type: ignore
 import pytest
 
 from tests.utils import wait_for_port_open
@@ -92,6 +93,30 @@ def test_mixin_ipc_service(mosec_service, http_client):
     assert resp.json() == "equal"
 
 
+@pytest.mark.parametrize(
+    "mosec_service, http_client",
+    [
+        pytest.param("mixin_typed_service", "", id="typed"),
+    ],
+    indirect=["mosec_service", "http_client"],
+)
+def test_mixin_typed_service(mosec_service, http_client):
+    resp = http_client.post(
+        "/inference",
+        content=msgpack.packb(
+            {
+                "media": "text",
+                "binary": b"hello mosec",
+            }
+        ),
+    )
+    assert resp.status_code == HTTPStatus.OK, resp
+    assert msgpack.unpackb(resp.content) == 11
+
+    resp = http_client.post("/inference", content=msgpack.packb({"media": "none"}))
+    assert resp.status_code == HTTPStatus.UNPROCESSABLE_ENTITY, resp
+
+
 @pytest.mark.parametrize(
     "mosec_service, http_client",
     [
diff --git a/tests/test_utils.py b/tests/test_utils.py
new file mode 100644
index 00000000..99e70421
--- /dev/null
+++ b/tests/test_utils.py
@@ -0,0 +1,42 @@
+# Copyright 2023 MOSEC Authors
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from typing import List
+
+from msgspec import Struct
+
+from mosec import Worker
+from mosec.mixin.typed_worker import parse_forward_input_type
+
+
+class Request(Struct):
+    name: str
+
+
+class Demo(Worker):
+    def forward(self, _data: Request):
+        pass
+
+    def batch_forward(self, _data: List[Request]):
+        pass
+
+
+def test_parse_forward_input_type():
+    demo = Demo()
+
+    single = parse_forward_input_type(demo.forward)
+    assert single is Request, single
+
+    batch = parse_forward_input_type(demo.batch_forward)
+    assert batch is Request, batch