Repository: mosecorg/mosec (language: python)

Issue:
[FEATURE] add numbin as a IPC serialization method
According to the benchmark, [numbin](https://github.com/mosecorg/numbin) should be better than pickle.

BTW, pickle is not secure to `loads`. (Should be okay for our use case, but not recommended)



Discussion / hints:

Fix patch:
diff --git a/.github/workflows/check.yml b/.github/workflows/check.yml
index dca8862d..2402f22b 100644
--- a/.github/workflows/check.yml
+++ b/.github/workflows/check.yml
@@ -73,11 +73,11 @@ jobs:
       - name: Install dependencies
         run: |
           python -m pip install --upgrade pip
-          pip install -e .[dev]
+          pip install -e .[dev,plugin]
           make dev
       - name: Test
         run: |
           make semantic_lint test
-      - name: Test plugins
-        if: ${{ startsWith(matrix.os, 'ubuntu') && matrix.python-version != '3.11' }}
-        run: make test_plugin
+      - name: Test pyarrow in Linux
+        if: ${{ startsWith(matrix.os, 'ubuntu') }}
+        run: make test_arrow
diff --git a/.github/workflows/package.yml b/.github/workflows/package.yml
index 65d7f7fc..248234d9 100644
--- a/.github/workflows/package.yml
+++ b/.github/workflows/package.yml
@@ -52,13 +52,13 @@ jobs:
       - name: Install dependencies
         run: |
           python -m pip install --upgrade pip
-          pip install -e .[dev]
+          pip install -e .[dev,plugin]
           make dev
       - name: Test
         run: make test
-      - name: Test plugins
-        if: ${{ startsWith(matrix.os, 'ubuntu') && matrix.python-version != '3.11' }}
-        run: make test_plugin
+      - name: Test pyarrow in Linux
+        if: ${{ startsWith(matrix.os, 'ubuntu') }}
+        run: make test_arrow
 
   build:
     name: "Build PyPI Package"
diff --git a/Makefile b/Makefile
index 0ab106d3..0c830959 100644
--- a/Makefile
+++ b/Makefile
@@ -20,7 +20,7 @@ test: dev
 	pytest tests -vv -s -m "not arrow"
 	RUST_BACKTRACE=1 cargo test -vv
 
-test_plugin: dev
+test_arrow: dev
 	@pip install -q -r requirements/plugin.txt
 	echo "Running tests for the plugin"
 	pytest tests -vv -s -m "arrow"
diff --git a/mosec/mixin/__init__.py b/mosec/mixin/__init__.py
index fc25af2b..043a17fc 100644
--- a/mosec/mixin/__init__.py
+++ b/mosec/mixin/__init__.py
@@ -15,5 +15,6 @@
 """Provide useful mixin to extend MOSEC."""
 
 from .msgpack_worker import MsgpackMixin
+from .numbin_worker import NumBinIPCMixin
 
-__all__ = ["MsgpackMixin"]
+__all__ = ["MsgpackMixin", "NumBinIPCMixin"]
diff --git a/mosec/mixin/numbin_worker.py b/mosec/mixin/numbin_worker.py
new file mode 100644
index 00000000..4af6b708
--- /dev/null
+++ b/mosec/mixin/numbin_worker.py
@@ -0,0 +1,55 @@
+# 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 NumBin IPC worker mixin.
+
+Features:
+
+    * deserialize IPC data with numbin
+    * serialize IPC data with numbin
+
+Attention: numbin only supports NumPy ndarray types.
+"""
+
+import warnings
+from typing import Any
+
+from ..errors import DecodingError, EncodingError
+
+try:
+    import numbin as nb  # type: ignore
+except ImportError:
+    warnings.warn("numbin is required for NumBinIPCMixin", ImportWarning)
+
+
+class NumBinIPCMixin:
+    """NumBin IPC worker mixin interface."""
+
+    # pylint: disable=no-self-use
+
+    def serialize_ipc(self, data: Any) -> bytes:
+        """Serialize with NumBin for the IPC."""
+        try:
+            data_bytes = nb.dumps(data)
+        except Exception as err:
+            raise EncodingError from err
+        return data_bytes
+
+    def deserialize_ipc(self, data: bytes) -> Any:
+        """Deserialize with NumBin for the IPC."""
+        try:
+            array = nb.loads(data)
+        except Exception as err:
+            raise DecodingError from err
+        return array
diff --git a/requirements/dev.txt b/requirements/dev.txt
index 1af39173..00571952 100644
--- a/requirements/dev.txt
+++ b/requirements/dev.txt
@@ -8,6 +8,5 @@ pydocstyle>=6.1.1
 black>=20.8.0
 isort>=5.6
 autoflake>=1.4
-msgpack>=1.0.2
 pre-commit>=2.15.0
 httpx==0.23.0
diff --git a/requirements/plugin.txt b/requirements/plugin.txt
index facba2e5..46acb22c 100644
--- a/requirements/plugin.txt
+++ b/requirements/plugin.txt
@@ -1,2 +1,3 @@
 pyarrow>=0.6.1
 msgpack>=1.0.2
+numbin>=0.4.0


Test patch:
diff --git a/tests/mixin_numbin_service.py b/tests/mixin_numbin_service.py
new file mode 100644
index 00000000..0538245a
--- /dev/null
+++ b/tests/mixin_numbin_service.py
@@ -0,0 +1,40 @@
+# 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 Dict, List
+
+import numpy as np
+
+from mosec import Server, Worker
+from mosec.mixin import NumBinIPCMixin
+
+
+class Preprocess(NumBinIPCMixin, Worker):
+    def forward(self, data: Dict[str, str]) -> np.ndarray:
+        num = int(data.get("num", 10))
+        arr = np.ones(num) * (1 / num)
+        return arr
+
+
+class Inference(NumBinIPCMixin, Worker):
+    def forward(self, data: List[np.ndarray]) -> List[str]:
+        res = ["equal" if np.equal(1, arr.sum()) else "unequal" for arr in data]
+        return res
+
+
+if __name__ == "__main__":
+    server = Server()
+    server.append_worker(Preprocess)
+    server.append_worker(Inference, max_batch_size=8)
+    server.run()
diff --git a/tests/test_coordinator.py b/tests/test_coordinator.py
index 401a7ba3..6fa6fad9 100644
--- a/tests/test_coordinator.py
+++ b/tests/test_coordinator.py
@@ -66,7 +66,7 @@ def forward(self, data):
         return data
 
 
-class EchoWorkerMSGPACK(MsgpackMixin, EchoWorkerJSON):
+class EchoWorkerMsgPack(MsgpackMixin, EchoWorkerJSON):
     """"""
 
 
@@ -209,7 +209,7 @@ def test_incorrect_socket_file(mocker, base_test_config, caplog):
                 msgpack.packb({"rid": "147982364", "data": b"im_bytes"}),
                 msgpack.packb({"rid": "147982831", "data": b"another_im_bytes"}),
             ],
-            EchoWorkerMSGPACK,
+            EchoWorkerMsgPack,
             msgpack.unpackb,
         ),
     ],
diff --git a/tests/test_service.py b/tests/test_service.py
index aff134e9..504b53de 100644
--- a/tests/test_service.py
+++ b/tests/test_service.py
@@ -17,13 +17,12 @@
 import shlex
 import subprocess
 import time
+from http import HTTPStatus
 from threading import Thread
 
 import httpx
 import pytest
 
-import mosec
-
 TEST_PORT = "5000"
 
 
@@ -61,24 +60,37 @@ def mosec_service(request):
 )
 def test_square_service(mosec_service, http_client):
     resp = http_client.get("/")
-    assert resp.status_code == 200
+    assert resp.status_code == HTTPStatus.OK
     # only check the prefix since the version from setuptools_scm may not be the
     # correct one used in `Cargo.toml`
     assert resp.headers["server"].startswith("mosec/"), f"{resp.headers['server']}"
 
     resp = http_client.get("/metrics")
-    assert resp.status_code == 200
+    assert resp.status_code == HTTPStatus.OK
 
     resp = http_client.post("/inference", json={"msg": 2})
-    assert resp.status_code == 422
+    assert resp.status_code == HTTPStatus.UNPROCESSABLE_ENTITY
     assert resp.text == "request validation error: 'x'"
 
     resp = http_client.post("/inference", content=b"bad-binary-request")
-    assert resp.status_code == 400
+    assert resp.status_code == HTTPStatus.BAD_REQUEST
 
     validate_square_service(http_client, 2)
 
 
+@pytest.mark.parametrize(
+    "mosec_service, http_client",
+    [
+        pytest.param(("mixin_numbin_service", 3), "", id="numbin"),
+    ],
+    indirect=["mosec_service", "http_client"],
+)
+def test_mixin_ipc_service(mosec_service, http_client):
+    resp = http_client.post("/inference", json={"num": 8})
+    assert resp.status_code == HTTPStatus.OK
+    assert resp.json() == "equal"
+
+
 @pytest.mark.parametrize(
     "mosec_service, http_client",
     [
@@ -109,6 +121,7 @@ def test_square_service_mp(mosec_service, http_client):
 
 def validate_square_service(http_client, x):
     resp = http_client.post("/inference", json={"x": x})
+    assert resp.status_code == HTTPStatus.OK
     assert resp.json()["x"] == x**2