pynumaflow-lite: Pythonic APIs for unary/batchmap/mapstream

.github/workflows/pynumaflow-lite.yml

@@ -30,6 +30,7 @@ env:
   RUSTFLAGS: -C debuginfo=0
   RUST_BACKTRACE: 1
   PYTHONUTF8: 1
+  RUST_TOOLCHAIN: "1.97"
 
 jobs:
   lint:
@@ -55,8 +56,8 @@ jobs:
 
       - name: Set up Rust
         run: |
-          rustup toolchain install stable --profile minimal --component clippy --component rustfmt
-          rustup default stable
+          rustup toolchain install "${RUST_TOOLCHAIN}" --profile minimal --component clippy --component rustfmt
+          rustup default "${RUST_TOOLCHAIN}"
 
       - name: Install dependencies
         run: uv sync --locked --group dev
@@ -100,8 +101,8 @@ jobs:
 
       - name: Set up Rust
         run: |
-          rustup toolchain install stable --profile minimal
-          rustup default stable
+          rustup toolchain install "${RUST_TOOLCHAIN}" --profile minimal
+          rustup default "${RUST_TOOLCHAIN}"
 
       - name: Run tests for Python versions
         shell: bash
@@ -173,8 +174,8 @@ jobs:
 
       - name: Set up Rust
         run: |
-          rustup toolchain install stable --profile minimal
-          rustup default stable
+          rustup toolchain install "${RUST_TOOLCHAIN}" --profile minimal
+          rustup default "${RUST_TOOLCHAIN}"
 
       - name: Verify tag matches package version
         if: github.event_name == 'push' && startsWith(github.ref, 'refs/tags/pynumaflow-lite-v')

packages/pynumaflow-lite/manifests/batchmap/batchmap_cat.py

@@ -1,46 +1,20 @@
-import asyncio
-import signal
-from collections.abc import AsyncIterable, Awaitable, Callable
+from collections.abc import AsyncIterator
 
 from pynumaflow_lite import batchmapper
-from pynumaflow_lite.batchmapper import Message
+from pynumaflow_lite.batchmapper import BatchResponse, Datum, Message
 
 
-class SimpleBatchCat(batchmapper.BatchMapper):
-    async def handler(self, batch: AsyncIterable[batchmapper.Datum]) -> batchmapper.BatchResponses:
-        responses = batchmapper.BatchResponses()
-        async for d in batch:
-            resp = batchmapper.BatchResponse(d.id)
-            if d.value == b"bad world":
-                resp.append(Message.message_to_drop())
-                continue
-
-            resp.append(Message(d.value, d.keys))
-            responses.append(resp)
-        return responses
-
-
-async def start(
-    f: Callable[[AsyncIterable[batchmapper.Datum]], Awaitable[batchmapper.BatchResponses]],
-):
-    server = batchmapper.BatchMapAsyncServer()
-
-    # Register loop-level signal handlers so we control shutdown and avoid asyncio.run
-    loop = asyncio.get_running_loop()
-    try:
-        loop.add_signal_handler(signal.SIGINT, lambda: server.stop())
-        loop.add_signal_handler(signal.SIGTERM, lambda: server.stop())
-    except (NotImplementedError, RuntimeError):
-        pass
-
-    try:
-        await server.start(f)
-        print("Shutting down gracefully...")
-    except asyncio.CancelledError:
-        server.stop()
-        return
+class SimpleBatchCat:
+    async def handler(self, batch: AsyncIterator[Datum]) -> list[BatchResponse]:
+        return [
+            BatchResponse(
+                d.id,
+                [Message.to_drop()] if d.value == b"bad world" else [Message(d.value, keys=d.keys)],
+            )
+            async for d in batch
+        ]
 
 
 if __name__ == "__main__":
-    async_handler = SimpleBatchCat()
-    asyncio.run(start(async_handler))
+    batch_mapper_obj = SimpleBatchCat()
+    batchmapper.BatchMapAsyncServer(batch_mapper_obj.handler).run()

packages/pynumaflow-lite/manifests/map/map_cat.py

@@ -1,48 +1,13 @@
-import asyncio
-import signal
-from collections.abc import Awaitable, Callable
-
 from pynumaflow_lite import mapper
 
 
-class SimpleCat(mapper.Mapper):
-    async def handler(self, keys: list[str], payload: mapper.Datum) -> mapper.Messages:
-
-        messages = mapper.Messages()
-
-        if payload.value == b"bad world":
-            messages.append(mapper.Message.message_to_drop())
-        else:
-            messages.append(mapper.Message(payload.value, keys))
-
-        return messages
-
-
-async def start(f: Callable[[list[str], mapper.Datum], Awaitable[mapper.Messages]]):
-    server = mapper.MapAsyncServer()
-
-    # Register loop-level signal handlers so we control shutdown and avoid asyncio.run
-    # converting it into KeyboardInterrupt/CancelledError traces.
-    loop = asyncio.get_running_loop()
-    loop.set_debug(True)
-    print("Registering signal handlers", loop)
-    try:
-        loop.add_signal_handler(signal.SIGINT, lambda: server.stop())
-        loop.add_signal_handler(signal.SIGTERM, lambda: server.stop())
-    except (NotImplementedError, RuntimeError):
-        print("Failed to register signal handlers")
-        # add_signal_handler may not be available on some platforms/contexts; fallback below.
-        pass
-
-    try:
-        await server.start(f)
-        print("Shutting down gracefully...")
-    except asyncio.CancelledError:
-        # Fallback in case the task was cancelled by the runner
-        server.stop()
-        return
+class SimpleCat:
+    async def handler(self, datum: mapper.Datum) -> list[mapper.Message]:
+        if datum.value == b"bad world":
+            return [mapper.Message.to_drop()]
+        return [mapper.Message(datum.value, keys=datum.keys)]
 
 
 if __name__ == "__main__":
-    async_handler = SimpleCat()
-    asyncio.run(start(async_handler))
+    mapper_obj = SimpleCat()
+    mapper.MapAsyncServer(mapper_obj.handler).run()

packages/pynumaflow-lite/manifests/mapstream/mapstream_cat.py

@@ -1,41 +1,19 @@
-import asyncio
-import signal
-from collections.abc import AsyncIterator, Callable
+from collections.abc import AsyncIterable
 
 from pynumaflow_lite import mapstreamer
 from pynumaflow_lite.mapstreamer import Message
 
 
-class SimpleStreamCat(mapstreamer.MapStreamer):
-    async def handler(self, keys: list[str], datum: mapstreamer.Datum) -> AsyncIterator[Message]:
+class SimpleStreamCat:
+    async def handler(self, datum: mapstreamer.Datum) -> AsyncIterable[Message]:
         parts = datum.value.decode("utf-8").split(",")
         if not parts:
             yield Message.to_drop()
             return
         for s in parts:
-            yield Message(s.encode(), keys)
-
-
-async def start(f: Callable[[list[str], mapstreamer.Datum], AsyncIterator[Message]]):
-    # Use default socket/info file locations; no explicit sock file passed
-    server = mapstreamer.MapStreamAsyncServer()
-
-    # Register loop-level signal handlers so we control shutdown and avoid asyncio.run noise.
-    loop = asyncio.get_running_loop()
-    try:
-        loop.add_signal_handler(signal.SIGINT, lambda: server.stop())
-        loop.add_signal_handler(signal.SIGTERM, lambda: server.stop())
-    except (NotImplementedError, RuntimeError):
-        pass
-
-    try:
-        await server.start(f)
-        print("Shutting down gracefully...")
-    except asyncio.CancelledError:
-        server.stop()
-        return
+            yield Message(s.encode(), keys=datum.keys)
 
 
 if __name__ == "__main__":
-    async_handler = SimpleStreamCat()
-    asyncio.run(start(async_handler))
+    map_streamer_obj = SimpleStreamCat()
+    mapstreamer.MapStreamAsyncServer(map_streamer_obj.handler).run()

packages/pynumaflow-lite/manifests/sideinput/sideinput_example.py

@@ -39,7 +39,7 @@ async def retrieve_handler(self) -> sideinputer.Response:
         return sideinputer.Response.broadcast_message(val.encode("utf-8"))
 
 
-class SideInputHandler(mapper.Mapper):
+class SideInputHandler:
     """
     A Mapper that reads from side input files and includes the value in its output.
     """
@@ -51,13 +51,11 @@ class SideInputHandler(mapper.Mapper):
     # Side input file that we are watching
     watched_file = "myticker"
 
-    async def handler(self, keys: list[str], datum: mapper.Datum) -> mapper.Messages:
+    async def handler(self, datum: mapper.Datum) -> list[mapper.Message]:
         with self.data_value_lock:
             current_value = self.data_value
 
-        messages = mapper.Messages()
-        messages.append(mapper.Message(str.encode(current_value)))
-        return messages
+        return [mapper.Message(str.encode(current_value))]
 
     def file_watcher(self):
         """
@@ -103,27 +101,18 @@ async def start_sideinput():
         server.stop()
 
 
-async def start_mapper():
+def start_mapper():
     """Start the Mapper server that reads from side inputs."""
-    server = mapper.MapAsyncServer()
-    handler = SideInputHandler()
+    mapper_obj = SideInputHandler()
 
     # Initialize the data value from the side input file
-    handler.init_data_value()
+    mapper_obj.init_data_value()
 
     # Start the file watcher in a background thread
-    watcher_thread = Thread(target=handler.file_watcher, daemon=True)
+    watcher_thread = Thread(target=mapper_obj.file_watcher, daemon=True)
     watcher_thread.start()
 
-    loop = asyncio.get_running_loop()
-    loop.add_signal_handler(signal.SIGINT, lambda: server.stop())
-    loop.add_signal_handler(signal.SIGTERM, lambda: server.stop())
-
-    try:
-        await server.start(handler)
-        print("Mapper server shutting down gracefully...")
-    except asyncio.CancelledError:
-        server.stop()
+    mapper.MapAsyncServer(mapper_obj.handler).run()
 
 
 if __name__ == "__main__":
@@ -132,7 +121,7 @@ async def start_mapper():
 
     if is_mapper:
         print("Starting as Mapper (reading side inputs)...")
-        asyncio.run(start_mapper())
+        start_mapper()
     else:
         print("Starting as SideInput retriever...")
         asyncio.run(start_sideinput())

packages/pynumaflow-lite/manifests/sink/sink_log.py

@@ -1,20 +1,19 @@
 import logging
-from collections.abc import AsyncIterable
+from collections.abc import AsyncIterator
 
 from pynumaflow_lite import sinker
-from pynumaflow_lite.sinker import Sinker
 
 # Configure logging
 logging.basicConfig(level=logging.INFO)
 _LOGGER = logging.getLogger(__name__)
 
 
-class SimpleLogSink(Sinker):
+class SimpleLogSink:
     """
     Simple log sink that logs each message and returns success responses.
     """
 
-    async def handler(self, datums: AsyncIterable[sinker.Datum]) -> list[sinker.Response]:
+    async def handler(self, datums: AsyncIterator[sinker.Datum]) -> list[sinker.Response]:
         responses = []
         async for msg in datums:
             _LOGGER.info("User Defined Sink: %s", msg.value.decode("utf-8"))
@@ -25,4 +24,5 @@ async def handler(self, datums: AsyncIterable[sinker.Datum]) -> list[sinker.Resp
 
 
 if __name__ == "__main__":
-    sinker.SinkAsyncServer(SimpleLogSink()).run()
+    sinker_obj = SimpleLogSink()
+    sinker.SinkAsyncServer(sinker_obj.handler).run()

packages/pynumaflow-lite/pynumaflow_lite/__init__.py

@@ -61,29 +61,27 @@
 except Exception:  # pragma: no cover
     sideinputer = None
 
-# Surface the Python Mapper, BatchMapper, MapStreamer, Reducer, SessionReducer, ReduceStreamer, Accumulator, Sinker,
-# Sourcer, SourceTransformer, and SideInput classes under the extension submodules for convenient access
+# Surface the Python async servers and data-type classes under the extension submodules for convenient access
 from ._accumulator_dtypes import Accumulator
-from ._batchmapper_dtypes import BatchMapper
-from ._map_dtypes import Mapper
-from ._mapstream_dtypes import MapStreamer
+from ._batchmap_server import BatchMapAsyncServer
+from ._map_server import MapAsyncServer
+from ._mapstream_server import MapStreamAsyncServer
 from ._reduce_dtypes import Reducer
 from ._reducestreamer_dtypes import ReduceStreamer
 from ._session_reduce_dtypes import SessionReducer
 from ._sideinput_dtypes import SideInput
-from ._sink_dtypes import Sinker
 from ._sink_server import SinkAsyncServer
 from ._source_dtypes import Sourcer
 from ._sourcetransformer_dtypes import SourceTransformer
 
 if mapper is not None:
-    mapper.Mapper = Mapper
+    mapper.MapAsyncServer = MapAsyncServer
 
 if batchmapper is not None:
-    batchmapper.BatchMapper = BatchMapper
+    batchmapper.BatchMapAsyncServer = BatchMapAsyncServer
 
 if mapstreamer is not None:
-    mapstreamer.MapStreamer = MapStreamer
+    mapstreamer.MapStreamAsyncServer = MapStreamAsyncServer
 
 if reducer is not None:
     reducer.Reducer = Reducer
@@ -98,7 +96,6 @@
     accumulator.Accumulator = Accumulator
 
 if sinker is not None:
-    sinker.Sinker = Sinker
     sinker.SinkAsyncServer = SinkAsyncServer
 
 if sourcer is not None: