feat: class based refactor for SDK

.codecov.yml

@@ -11,10 +11,4 @@ coverage:
 
 ignore:
   - "examples/"
-  - "pynumaflow/mapper/proto/*"
-  - "pynumaflow/sinker/proto/*"
-  - "pynumaflow/mapstreamer/proto/*"
-  - "pynumaflow/reducer/proto/*"
-  - "pynumaflow/sourcetransformer/proto/*"
-  - "pynumaflow/sideinput/proto/*"
-  - "pynumaflow/sourcer/proto/*"
+  - "pynumaflow/proto/*"

.coveragerc

@@ -5,14 +5,18 @@ source = pynumaflow
 omit =
     pynumaflow/tests/*
     examples/*
+    pynumaflow/proto/*
+    pynumaflow/shared/server.py
 
 [report]
 exclude_lines =
-    def start
-    def start_async
-    def __serve_async
-    def start_multiproc
+    def sync_server_start
     def _run_server
+    def start_multiproc_server
+    async def start_async_server
     def _reserve_port
     if os.getenv("PYTHONDEBUG"):
     _LOGGER.setLevel(logging.DEBUG)
+    def exec_multiproc
+    def exec
+    async def aexec

Makefile

@@ -27,13 +27,13 @@ setup:
 
 
 proto:
-	python3 -m grpc_tools.protoc -I=pynumaflow/sinker/proto --python_out=pynumaflow/sinker/proto --grpc_python_out=pynumaflow/sinker/proto  pynumaflow/sinker/proto/*.proto
-	python3 -m grpc_tools.protoc -I=pynumaflow/mapper/proto --python_out=pynumaflow/mapper/proto --grpc_python_out=pynumaflow/mapper/proto  pynumaflow/mapper/proto/*.proto
-	python3 -m grpc_tools.protoc -I=pynumaflow/mapstreamer/proto --python_out=pynumaflow/mapstreamer/proto --grpc_python_out=pynumaflow/mapstreamer/proto  pynumaflow/mapstreamer/proto/*.proto
-	python3 -m grpc_tools.protoc -I=pynumaflow/reducer/proto --python_out=pynumaflow/reducer/proto --grpc_python_out=pynumaflow/reducer/proto  pynumaflow/reducer/proto/*.proto
-	python3 -m grpc_tools.protoc -I=pynumaflow/sourcetransformer/proto --python_out=pynumaflow/sourcetransformer/proto --grpc_python_out=pynumaflow/sourcetransformer/proto  pynumaflow/sourcetransformer/proto/*.proto
-	python3 -m grpc_tools.protoc -I=pynumaflow/sideinput/proto --python_out=pynumaflow/sideinput/proto --grpc_python_out=pynumaflow/sideinput/proto  pynumaflow/sideinput/proto/*.proto
-	python3 -m grpc_tools.protoc -I=pynumaflow/sourcer/proto --python_out=pynumaflow/sourcer/proto --grpc_python_out=pynumaflow/sourcer/proto  pynumaflow/sourcer/proto/*.proto
+	python3 -m grpc_tools.protoc -I=pynumaflow/proto/sinker --python_out=pynumaflow/proto/sinker --grpc_python_out=pynumaflow/proto/sinker  pynumaflow/proto/sinker/*.proto
+	python3 -m grpc_tools.protoc -I=pynumaflow/proto/mapper --python_out=pynumaflow/proto/mapper --grpc_python_out=pynumaflow/proto/mapper  pynumaflow/proto/mapper/*.proto
+	python3 -m grpc_tools.protoc -I=pynumaflow/proto/mapstreamer --python_out=pynumaflow/proto/mapstreamer --grpc_python_out=pynumaflow/proto/mapstreamer  pynumaflow/proto/mapstreamer/*.proto
+	python3 -m grpc_tools.protoc -I=pynumaflow/proto/reducer --python_out=pynumaflow/proto/reducer --grpc_python_out=pynumaflow/proto/reducer  pynumaflow/proto/reducer/*.proto
+	python3 -m grpc_tools.protoc -I=pynumaflow/proto/sourcetransformer --python_out=pynumaflow/proto/sourcetransformer --grpc_python_out=pynumaflow/proto/sourcetransformer  pynumaflow/proto/sourcetransformer/*.proto
+	python3 -m grpc_tools.protoc -I=pynumaflow/proto/sideinput --python_out=pynumaflow/proto/sideinput --grpc_python_out=pynumaflow/proto/sideinput  pynumaflow/proto/sideinput/*.proto
+	python3 -m grpc_tools.protoc -I=pynumaflow/proto/sourcer --python_out=pynumaflow/proto/sourcer --grpc_python_out=pynumaflow/proto/sourcer  pynumaflow/proto/sourcer/*.proto
 
 
-	sed -i '' 's/^\(import.*_pb2\)/from . \1/' pynumaflow/*/proto/*.py
+	sed -i '' 's/^\(import.*_pb2\)/from . \1/' pynumaflow/proto/*/*.py

README.md

@@ -46,18 +46,25 @@ pre-commit install
 ### Map
 
 ```python
-from pynumaflow.mapper import Messages, Message, Datum, Mapper
+from pynumaflow.mapper import Messages, Message, Datum, MapServer
 
 
-def my_handler(keys: list[str], datum: Datum) -> Messages:
+def handler(keys: list[str], datum: Datum) -> Messages:
     val = datum.value
     _ = datum.event_time
     _ = datum.watermark
-    return Messages(Message(value=val, keys=keys))
+    strs = val.decode("utf-8").split(",")
+    messages = Messages()
+    if len(strs) == 0:
+        messages.append(Message.to_drop())
+        return messages
+    for s in strs:
+        messages.append(Message(str.encode(s)))
+    return messages
 
 
 if __name__ == "__main__":
-    grpc_server = Mapper(handler=my_handler)
+    grpc_server = MapServer(handler)
     grpc_server.start()
 ```
 ### SourceTransformer - Map with event time assignment capability
@@ -66,7 +73,7 @@ SourceTransformer is only supported at source vertex to enable (a) early data fi
 
 ```python
 from datetime import datetime
-from pynumaflow.sourcetransformer import Messages, Message, Datum, SourceTransformer
+from pynumaflow.sourcetransformer import Messages, Message, Datum, SourceTransformServer
 
 
 def transform_handler(keys: list[str], datum: Datum) -> Messages:
@@ -78,21 +85,18 @@ def transform_handler(keys: list[str], datum: Datum) -> Messages:
 
 
 if __name__ == "__main__":
-    grpc_server = SourceTransformer(handler=transform_handler)
+    grpc_server = SourceTransformServer(transform_handler)
     grpc_server.start()
 ```
 
 ### Reduce
 
 ```python
-import aiorun
-from typing import Iterator, List
-from pynumaflow.reducer import Messages, Message, Datum, Metadata, AsyncReducer
+from typing import AsyncIterable
+from pynumaflow.reducer import Messages, Message, Datum, Metadata, ReduceAsyncServer
 
 
-async def my_handler(
-        keys: List[str], datums: Iterator[Datum], md: Metadata
-) -> Messages:
+async def reduce_handler(keys: list[str], datums: AsyncIterable[Datum], md: Metadata) -> Messages:
     interval_window = md.interval_window
     counter = 0
     async for _ in datums:
@@ -101,12 +105,12 @@ async def my_handler(
         f"counter:{counter} interval_window_start:{interval_window.start} "
         f"interval_window_end:{interval_window.end}"
     )
-    return Messages(Message(str.encode(msg), keys))
+    return Messages(Message(str.encode(msg), keys=keys))
 
 
 if __name__ == "__main__":
-    grpc_server = AsyncReducer(handler=my_handler)
-    aiorun.run(grpc_server.start())
+    grpc_server = ReduceAsyncServer(reduce_handler)
+    grpc_server.start()
 ```
 
 ### Sample Image
@@ -117,7 +121,7 @@ under [examples](examples/map/forward_message).
 
 ```python
 from typing import Iterator
-from pynumaflow.sinker import Datum, Responses, Response, Sinker
+from pynumaflow.sinker import Datum, Responses, Response, SinkServer
 
 
 def my_handler(datums: Iterator[Datum]) -> Responses:
@@ -129,11 +133,84 @@ def my_handler(datums: Iterator[Datum]) -> Responses:
 
 
 if __name__ == "__main__":
-    grpc_server = Sinker(my_handler)
+    grpc_server = SinkServer(my_handler)
     grpc_server.start()
 ```
 
 ### Sample Image
 
 A sample UDSink [Dockerfile](examples/sink/log/Dockerfile) is provided
-under [examples](examples/sink/log).
+under [examples](examples/sink/log).
+
+## Class based handlers 
+
+We can also implement UDFs and UDSinks using class based handlers.
+
+The class based handlers are useful when we want to maintain state across multiple invocations of the handler.
+
+Here we can pass the class instance to the server and the server will invoke the handler methods on the instance.
+
+To use a class based handler, we the user needs to inherit the base class of the UDF/UDSink.
+And implement the required methods in the class.
+
+Example For Mapper, the user needs to inherit the [Mapper](pynumaflow/mapper/_dtypes.py#170) class and then implement the [handler](pynumaflow/mapper/_dtypes.py#170) method.
+
+### Map
+
+```python
+from pynumaflow.mapper import Messages, Message, Datum, MapServer, Mapper
+
+
+class MyHandler(Mapper):
+    def handler(self, keys: list[str], datum: Datum) -> Messages:
+        val = datum.value
+        _ = datum.event_time
+        _ = datum.watermark
+        strs = val.decode("utf-8").split(",")
+        messages = Messages()
+        if len(strs) == 0:
+            messages.append(Message.to_drop())
+            return messages
+        for s in strs:
+            messages.append(Message(str.encode(s)))
+        return messages
+
+
+if __name__ == "__main__":
+    class_instance = MyHandler()
+    grpc_server = MapServer(class_instance)
+    grpc_server.start()
+```
+
+
+## Server Types
+
+For different types of UDFs and UDSinks, we have different server types which are supported.
+
+These have different functionalities and are used for different use cases.
+
+Currently we support the following server types:
+1) SyncServer
+2) AsyncServer
+3) MultiProcessServer
+
+Not all of the above are supported for all UDFs and UDSinks.
+
+
+
+### SyncServer
+```
+grpc_server = MapServer(handler)
+```
+
+### AsyncServer
+```
+grpc_server = MapAsyncServer(handler)
+```
+
+### MultiProcessServer
+```
+grpc_server = MapMultiProcServer(handler)
+```
+
+

examples/developer_guide/example.py

@@ -1,29 +1,27 @@
-import aiorun
-from collections.abc import Iterator
+from collections.abc import AsyncIterable
+
 
 from pynumaflow.reducer import (
     Messages,
     Message,
     Datum,
     Metadata,
-    AsyncReducer,
+    ReduceAsyncServer,
 )
 
 
-async def my_handler(keys: list[str], datums: Iterator[Datum], md: Metadata) -> Messages:
-    # count the number of events
+async def reduce_handler(keys: list[str], datums: AsyncIterable[Datum], md: Metadata) -> Messages:
     interval_window = md.interval_window
     counter = 0
     async for _ in datums:
         counter += 1
-
     msg = (
         f"counter:{counter} interval_window_start:{interval_window.start} "
         f"interval_window_end:{interval_window.end}"
     )
-    return Messages(Message(keys=keys, value=str.encode(msg)))
+    return Messages(Message(str.encode(msg), keys=keys))
 
 
 if __name__ == "__main__":
-    grpc_server = AsyncReducer(handler=my_handler)
-    aiorun.run(grpc_server.start())
+    grpc_server = ReduceAsyncServer(reduce_handler)
+    grpc_server.start()

examples/map/even_odd/example.py

@@ -1,4 +1,4 @@
-from pynumaflow.mapper import Messages, Message, Datum, Mapper
+from pynumaflow.mapper import Messages, Message, Datum, MapServer
 
 
 def my_handler(keys: list[str], datum: Datum) -> Messages:
@@ -22,5 +22,5 @@ def my_handler(keys: list[str], datum: Datum) -> Messages:
 
 
 if __name__ == "__main__":
-    grpc_server = Mapper(handler=my_handler)
+    grpc_server = MapServer(my_handler)
     grpc_server.start()

examples/map/even_odd/pyproject.toml

@@ -6,7 +6,7 @@ authors = ["Numaflow developers"]
 
 [tool.poetry.dependencies]
 python = "~3.10"
-pynumaflow = "~0.6.0"
+pynumaflow = "~0.7.0"
 
 [tool.poetry.dev-dependencies]
 

examples/map/flatmap/example.py

@@ -1,20 +1,21 @@
-from pynumaflow.mapper import Messages, Message, Datum, Mapper
+from pynumaflow.mapper import Messages, Message, Datum, MapServer, Mapper
 
 
-def my_handler(keys: list[str], datum: Datum) -> Messages:
-    val = datum.value
-    _ = datum.event_time
-    _ = datum.watermark
-    strs = val.decode("utf-8").split(",")
-    messages = Messages()
-    if len(strs) == 0:
-        messages.append(Message.to_drop())
+class Flatmap(Mapper):
+    def handler(self, keys: list[str], datum: Datum) -> Messages:
+        val = datum.value
+        _ = datum.event_time
+        _ = datum.watermark
+        strs = val.decode("utf-8").split(",")
+        messages = Messages()
+        if len(strs) == 0:
+            messages.append(Message.to_drop())
+            return messages
+        for s in strs:
+            messages.append(Message(str.encode(s)))
         return messages
-    for s in strs:
-        messages.append(Message(str.encode(s)))
-    return messages
 
 
 if __name__ == "__main__":
-    grpc_server = Mapper(handler=my_handler)
+    grpc_server = MapServer(Flatmap())
     grpc_server.start()

examples/map/flatmap/pipeline.yaml

@@ -15,7 +15,7 @@ spec:
     - name: flatmap
       udf:
         container:
-          image: "quay.io/numaio/numaflow-python/map-flatmap:v0.5.0"
+          image: "quay.io/numaio/numaflow-python/map-flatmap:v0.7.0"
           env:
             - name: PYTHONDEBUG
               value: "true"

examples/map/flatmap/pyproject.toml

@@ -6,7 +6,7 @@ authors = ["Numaflow developers"]
 
 [tool.poetry.dependencies]
 python = "~3.10"
-pynumaflow = "~0.6.0"
+pynumaflow = "~0.7.0"
 
 
 [tool.poetry.dev-dependencies]

examples/map/forward_message/Makefile

@@ -1,6 +1,6 @@
 .PHONY: image
 image:
-	docker build -t "quay.io/numaio/numaflow-python/map-forward-message:v0.5.0" .
+	docker build -t "quay.io/numaio/numaflow-python/map-forward-message:v0.7.0" .
 # Github CI runner uses platform linux/amd64. If your local environment don't, the image built by command above might not work
 # under the CI E2E test environment.
 # To build an image that supports multiple platforms(linux/amd64,linux/arm64) and push to quay.io, use the following command

examples/map/forward_message/example.py

@@ -1,4 +1,16 @@
-from pynumaflow.mapper import Messages, Message, Datum, Mapper
+import os
+
+from pynumaflow.mapper import Messages, Message, Datum, MapServer, Mapper
+
+
+class Example(Mapper):
+    def handler(self, keys: list[str], datum: Datum) -> Messages:
+        val = datum.value
+        _ = datum.event_time
+        _ = datum.watermark
+        messages = Messages()
+        messages.append(Message(value=val, keys=keys))
+        return messages
 
 
 def my_handler(keys: list[str], datum: Datum) -> Messages:
@@ -11,5 +23,13 @@ def my_handler(keys: list[str], datum: Datum) -> Messages:
 
 
 if __name__ == "__main__":
-    grpc_server = Mapper(handler=my_handler)
+    invoke = os.getenv("INVOKE", "handler")
+    # Use the class based approach or function based handler
+    # based on the env variable
+    # Both can be used and passed directly to the server class
+    if invoke == "class":
+        handler = Example()
+    else:
+        handler = my_handler
+    grpc_server = MapServer(handler)
     grpc_server.start()

examples/map/forward_message/pyproject.toml

@@ -6,7 +6,7 @@ authors = ["Numaflow developers"]
 
 [tool.poetry.dependencies]
 python = "~3.10"
-pynumaflow = "~0.6.0"
+pynumaflow = "~0.7.0"
 
 [tool.poetry.dev-dependencies]