Repository: mosecorg/mosec (language: python)

Issue:
feat: provide redis shared memory mixin
### Describe the feature

Provide redis as ipc wrapper plugin for better performance than plasma. And in some scenarios, such as stateful inference, users may need redis as distributed consistent store.

### Why do you need this feature?

_No response_

### Additional context

_No response_


Discussion / hints:

Fix patch:
diff --git a/.github/workflows/check.yml b/.github/workflows/check.yml
index 8b21cd31..9807872d 100644
--- a/.github/workflows/check.yml
+++ b/.github/workflows/check.yml
@@ -81,4 +81,6 @@ jobs:
           make semantic_lint test
       - name: Test pyarrow in Linux
         if: ${{ startsWith(matrix.os, 'ubuntu') }}
-        run: make test_arrow
+        run: |
+          sudo apt update && sudo apt install redis
+          make test_arrow
diff --git a/docs/source/examples/ipc.md b/docs/source/examples/ipc.md
index 9d9eb7be..36e08d52 100644
--- a/docs/source/examples/ipc.md
+++ b/docs/source/examples/ipc.md
@@ -4,7 +4,7 @@ This is an example demonstrating how you can enable the plasma shared memory sto
 
 Mosec's multi-stage pipeline requires the output data from the previous stage to be transferred to the next stage across python processes. This is coordinated via Unix domain socket between every Python worker process from all stages and the Rust controller process.
 
-By default, we serialize the data and directly transfer the bytes over the socket. However, users may find wrapping this IPC useful or more efficient for specific use cases. Therefore, we provide an example implementation `PlasmaShmIPCMixin` based on [`pyarrow.plasma`](https://arrow.apache.org/docs/11.0/python/plasma.html).
+By default, we serialize the data and directly transfer the bytes over the socket. However, users may find wrapping this IPC useful or more efficient for specific use cases. Therefore, we provide an example implementation `PlasmaShmIPCMixin` based on [`pyarrow.plasma`](https://arrow.apache.org/docs/11.0/python/plasma.html) and `RedisShmIPCMixin` based on [`redis`](https://pypi.org/project/redis). We recommand using `RedisShmWrapper` for better performance and longer-lasting updates.
 
 ```{warning}
 `plasma` is deprecated. Please use Redis instead.
@@ -17,6 +17,11 @@ The additional subprocess can be registered as a daemon thus it will be checked
 ```{include} ../../../examples/plasma_shm_ipc.py
 :code: python
 ```
+## **`redis_shm_ipc.py`**
+
+```{include} ../../../examples/redis_shm_ipc.py
+:code: python
+```
 
 ## Start
 
@@ -24,6 +29,11 @@ The additional subprocess can be registered as a daemon thus it will be checked
 python plasma_shm_ipc.py
 ```
 
+or
+
+```shell
+python redis_shm_ipc.py
+```
 ## Test
 
 ```shell
diff --git a/examples/redis_shm_ipc.py b/examples/redis_shm_ipc.py
new file mode 100644
index 00000000..7470c0a4
--- /dev/null
+++ b/examples/redis_shm_ipc.py
@@ -0,0 +1,60 @@
+# 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.
+
+"""Example: Using Redis store with mosec mixin RedisShmIPCMixin.
+
+We start a subprocess for the Redis server, and pass the url
+to the redis client which serves as the shm mixin.
+We also register the redis server process as a daemon, so
+that when it exits the service is able to gracefully shut down
+and be restarted by the orchestrator.
+"""
+
+import subprocess
+
+import numpy as np
+
+from mosec import Server, ValidationError, Worker
+from mosec.mixin import RedisShmIPCMixin
+
+
+class DataProducer(RedisShmIPCMixin, Worker):
+    """Sample Data Producer."""
+
+    def forward(self, data: dict) -> np.ndarray:
+        try:
+            nums = np.random.rand(int(data["size"]))
+        except KeyError as err:
+            raise ValidationError(err) from err
+        return nums
+
+
+class DataConsumer(RedisShmIPCMixin, Worker):
+    """Sample Data Consumer."""
+
+    def forward(self, data: np.ndarray) -> dict:
+        return {"ipc test data": data.tolist()}
+
+
+if __name__ == "__main__":
+    with subprocess.Popen(["redis-server"]) as p:  # start the redis server
+        # configure the redis url
+        RedisShmIPCMixin.set_redis_url("redis://localhost:6379/0")
+
+        server = Server()
+        # register this process to be monitored
+        server.register_daemon("redis-server", p)
+        server.append_worker(DataProducer, num=2)
+        server.append_worker(DataConsumer, num=2)
+        server.run()
diff --git a/mosec/mixin/__init__.py b/mosec/mixin/__init__.py
index 4479a99f..54b083fb 100644
--- a/mosec/mixin/__init__.py
+++ b/mosec/mixin/__init__.py
@@ -17,6 +17,13 @@
 from mosec.mixin.msgpack_worker import MsgpackMixin
 from mosec.mixin.numbin_worker import NumBinIPCMixin
 from mosec.mixin.plasma_worker import PlasmaShmIPCMixin
+from mosec.mixin.redis_worker import RedisShmIPCMixin
 from mosec.mixin.typed_worker import TypedMsgPackMixin
 
-__all__ = ["MsgpackMixin", "NumBinIPCMixin", "PlasmaShmIPCMixin", "TypedMsgPackMixin"]
+__all__ = [
+    "MsgpackMixin",
+    "NumBinIPCMixin",
+    "PlasmaShmIPCMixin",
+    "TypedMsgPackMixin",
+    "RedisShmIPCMixin",
+]
diff --git a/mosec/mixin/redis_worker.py b/mosec/mixin/redis_worker.py
new file mode 100644
index 00000000..2a1d4f2b
--- /dev/null
+++ b/mosec/mixin/redis_worker.py
@@ -0,0 +1,97 @@
+# 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 redis worker mixin.
+
+Provide another data transfer way between workers.
+
+The data will be stored in redis shared memory, while the object ID will be
+sent via the original way.
+
+    use case: large image tensors, cluster-shared data
+    benefits: more stable P99 latency
+
+"""
+import warnings
+from os import environ
+from typing import Any
+
+try:
+    import redis  # type: ignore
+except ImportError:
+    warnings.warn(
+        "redis is not installed. RedisShmIPCMixin is not available.", ImportWarning
+    )
+
+
+from mosec.worker import Worker
+
+_REDIS_URL_ENV = "MOSEC_INTERNAL_REDIS_URL"
+_DEFAULT_KEY = "REDIS_SHM_IPC_KEY"
+
+
+class RedisShmIPCMixin(Worker):
+    """Redis shared memory worker mixin interface."""
+
+    # pylint: disable=no-self-use
+
+    _redis_client = None
+    _redis_key = _DEFAULT_KEY
+    _next_id = None
+
+    @classmethod
+    def set_redis_url(cls, url: str):
+        """Set the redis service url."""
+        environ[_REDIS_URL_ENV] = url
+
+    def _get_client(self):
+        """Get the redis client. This will create a new one if not exist."""
+        if not self._redis_client:
+            url = environ.get(_REDIS_URL_ENV)
+            if not url:
+                raise RuntimeError(
+                    "please set the redis url with "
+                    "`RedisShmIPCMixin.set_redis_url()`"
+                )
+            self._redis_client = redis.from_url(url)
+        return self._redis_client
+
+    def _prepare_next_id(self) -> None:
+        """Make sure the next id exists. This will create a new one if not exist."""
+        if self._next_id is None:
+            client = self._get_client()
+            key = self._redis_key
+            self._next_id = bytes(str(client.incr(key)), encoding="utf-8")
+
+    def serialize_ipc(self, data: Any) -> bytes:
+        """Save the data to the redis server and return the id."""
+        self._prepare_next_id()
+        client = self._get_client()
+        with client.pipeline() as pipe:
+            current_id = self._next_id
+            pipe.set(current_id, super().serialize_ipc(data))  # type: ignore
+            pipe.incr(self._redis_key)
+            _id = pipe.execute()[-1]
+            self._next_id = bytes(str(_id), encoding="utf-8")
+        return current_id  # type: ignore
+
+    def deserialize_ipc(self, data: bytes) -> Any:
+        """Get the data from the redis server and delete it."""
+        client = self._get_client()
+        object_id = bytes(data)
+        with client.pipeline() as pipe:
+            pipe.get(object_id)
+            pipe.delete(object_id)
+            obj = pipe.execute()[0]
+        return super().deserialize_ipc(obj)
diff --git a/requirements/mixin.txt b/requirements/mixin.txt
index ba817f7d..1efe4fab 100644
--- a/requirements/mixin.txt
+++ b/requirements/mixin.txt
@@ -2,3 +2,4 @@ msgpack>=1.0.2
 msgspec>=0.15.0
 numbin>=0.4.0
 pyarrow>=0.6.1,<12
+redis>=4.0.0
\ No newline at end of file


Test patch:
diff --git a/tests/mixin_ipc_shm_service.py b/tests/mixin_ipc_shm_service.py
index 812fc76c..0aee1068 100644
--- a/tests/mixin_ipc_shm_service.py
+++ b/tests/mixin_ipc_shm_service.py
@@ -12,6 +12,8 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+import subprocess
+import sys
 from typing import List
 
 import numpy as np
@@ -19,10 +21,10 @@
 
 from mosec import Server, Worker
 from mosec.errors import ValidationError
-from mosec.mixin import PlasmaShmIPCMixin
+from mosec.mixin import PlasmaShmIPCMixin, RedisShmIPCMixin
 
 
-class RandomService(PlasmaShmIPCMixin, Worker):
+class PlasmaRandomService(PlasmaShmIPCMixin, Worker):
     def forward(self, data: List[dict]) -> List[dict]:
         try:
             result = [{"x": np.random.rand(int(req["size"]))} for req in data]
@@ -31,7 +33,7 @@ def forward(self, data: List[dict]) -> List[dict]:
         return result
 
 
-class DummyPostprocess(PlasmaShmIPCMixin, Worker):
+class PlasmaDummyPostprocess(PlasmaShmIPCMixin, Worker):
     """This dummy stage is added to test the shm IPC"""
 
     def forward(self, data: dict) -> dict:
@@ -40,7 +42,38 @@ def forward(self, data: dict) -> dict:
         return data
 
 
-if __name__ == "__main__":
+class RedisRandomService(RedisShmIPCMixin, Worker):
+    def forward(self, data: List[dict]) -> List[dict]:
+        try:
+            result = [{"x": np.random.rand(int(req["size"]))} for req in data]
+        except KeyError as err:
+            raise ValidationError(err)
+        return result
+
+
+class RedisDummyPostprocess(RedisShmIPCMixin, Worker):
+    """This dummy stage is added to test the shm IPC"""
+
+    def forward(self, data: dict) -> dict:
+        assert isinstance(data.get("x"), np.ndarray), f"wrong data type: {data}"
+        data["x"] = data["x"].tolist()
+        return data
+
+
+def start_redis_shm_mosec():
+    with subprocess.Popen(["redis-server"]) as p:  # start the redis server
+        # configure the plasma service path
+        RedisShmIPCMixin.set_redis_url("redis://localhost:6379/0")
+
+        server = Server()
+        # register this process to be monitored
+        server.register_daemon("redis_server", p)
+        server.append_worker(RedisRandomService, max_batch_size=8)
+        server.append_worker(RedisDummyPostprocess, num=2)
+        server.run()
+
+
+def start_plasma_shm_mosec():
     # initialize a 20Mb object store as shared memory
     with plasma.start_plasma_store(plasma_store_memory=20 * 1000 * 1000) as (
         shm_path,
@@ -51,6 +84,18 @@ def forward(self, data: dict) -> dict:
 
         server = Server()
         server.register_daemon("plasma_server", shm_process)
-        server.append_worker(RandomService, max_batch_size=8)
-        server.append_worker(DummyPostprocess, num=2)
+        server.append_worker(PlasmaRandomService, max_batch_size=8)
+        server.append_worker(PlasmaDummyPostprocess, num=2)
         server.run()
+
+
+if __name__ == "__main__":
+    if len(sys.argv) < 2:
+        print("Please specify a shm to run: plasma or redis")
+        exit(1)
+
+    shm = sys.argv[1]
+    if shm == "plasma":
+        start_plasma_shm_mosec()
+    elif shm == "redis":
+        start_redis_shm_mosec()
diff --git a/tests/test_service.py b/tests/test_service.py
index 40f11f58..2ed72caa 100644
--- a/tests/test_service.py
+++ b/tests/test_service.py
@@ -37,9 +37,11 @@ def http_client():
 
 @pytest.fixture(scope="session")
 def mosec_service(request):
-    name = request.param
+    params = request.param.split(" ")
+    name = params[0]
+    args = " ".join(params[1:])
     service = subprocess.Popen(
-        shlex.split(f"python -u tests/{name}.py --port {TEST_PORT}"),
+        shlex.split(f"python -u tests/{name}.py {args} --port {TEST_PORT}"),
     )
     assert wait_for_port_open(port=TEST_PORT), "service failed to start"
     yield service
@@ -78,7 +80,10 @@ def test_square_service(mosec_service, http_client):
     "mosec_service, http_client",
     [
         pytest.param(
-            "mixin_ipc_shm_service", "", id="shm_plasma", marks=pytest.mark.arrow
+            "mixin_ipc_shm_service plasma", "", id="shm_plasma", marks=pytest.mark.arrow
+        ),
+        pytest.param(
+            "mixin_ipc_shm_service redis", "", id="shm_redis", marks=pytest.mark.arrow
         ),
     ],
     indirect=["mosec_service", "http_client"],