From 9f9553e6e9315f6a26801ae153e0126b0ad5f071 Mon Sep 17 00:00:00 2001 From: Vigith Maurice Date: Thu, 10 Sep 2026 07:32:10 -0700 Subject: [PATCH 1/2] =?UTF-8?q?Revert=20"chore(pynumaflow-lite):=20Pythoni?= =?UTF-8?q?c=20APIs=20for=20unary/batchmap/mapstream=20(#=E2=80=A6"?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit This reverts commit 0bccb8dabe94c06305e204ab4ea62888d43b5720. --- .github/workflows/pynumaflow-lite.yml | 13 +- .../manifests/batchmap/batchmap_cat.py | 52 +- .../pynumaflow-lite/manifests/map/map_cat.py | 49 +- .../manifests/mapstream/mapstream_cat.py | 34 +- .../manifests/sideinput/sideinput_example.py | 29 +- .../manifests/sink/sink_log.py | 10 +- .../pynumaflow_lite/__init__.py | 17 +- .../pynumaflow_lite/_batchmap_server.py | 91 ---- .../pynumaflow_lite/_batchmapper_dtypes.py | 21 + .../pynumaflow_lite/_map_dtypes.py | 24 + .../pynumaflow_lite/_map_server.py | 88 ---- .../pynumaflow_lite/_mapstream_dtypes.py | 22 + .../pynumaflow_lite/_mapstream_server.py | 88 ---- .../pynumaflow_lite/_sink_dtypes.py | 21 + .../pynumaflow_lite/_sink_server.py | 19 +- .../pynumaflow_lite/batchmapper.pyi | 60 +-- .../pynumaflow_lite/mapper.pyi | 122 +++-- .../pynumaflow_lite/mapstreamer.pyi | 50 +- .../pynumaflow_lite/sinker.pyi | 14 +- packages/pynumaflow-lite/src/batchmap/mod.rs | 202 ++++---- .../pynumaflow-lite/src/batchmap/server.rs | 109 +---- packages/pynumaflow-lite/src/map/mod.rs | 463 ++++++++++-------- packages/pynumaflow-lite/src/map/server.rs | 119 +---- packages/pynumaflow-lite/src/mapstream/mod.rs | 131 ++--- .../pynumaflow-lite/src/mapstream/server.rs | 92 +--- packages/pynumaflow-lite/src/sink/server.rs | 2 +- packages/pynumaflow-lite/tests/_test_utils.py | 12 - .../tests/examples/batchmap_cat.py | 53 +- .../tests/examples/batchmap_cat_class.py | 58 ++- .../pynumaflow-lite/tests/examples/map_cat.py | 71 ++- .../tests/examples/map_cat_class.py | 82 +++- .../tests/examples/mapstream_cat.py | 35 +- .../tests/examples/mapstream_cat_class.py | 39 +- .../tests/examples/sink_log_class.py | 10 +- .../pynumaflow-lite/tests/test_map_unit.py | 92 ---- 35 files changed, 1041 insertions(+), 1353 deletions(-) delete mode 100644 packages/pynumaflow-lite/pynumaflow_lite/_batchmap_server.py create mode 100644 packages/pynumaflow-lite/pynumaflow_lite/_batchmapper_dtypes.py create mode 100644 packages/pynumaflow-lite/pynumaflow_lite/_map_dtypes.py delete mode 100644 packages/pynumaflow-lite/pynumaflow_lite/_map_server.py create mode 100644 packages/pynumaflow-lite/pynumaflow_lite/_mapstream_dtypes.py delete mode 100644 packages/pynumaflow-lite/pynumaflow_lite/_mapstream_server.py create mode 100644 packages/pynumaflow-lite/pynumaflow_lite/_sink_dtypes.py delete mode 100644 packages/pynumaflow-lite/tests/test_map_unit.py diff --git a/.github/workflows/pynumaflow-lite.yml b/.github/workflows/pynumaflow-lite.yml index feb79796..2d239831 100644 --- a/.github/workflows/pynumaflow-lite.yml +++ b/.github/workflows/pynumaflow-lite.yml @@ -30,7 +30,6 @@ env: RUSTFLAGS: -C debuginfo=0 RUST_BACKTRACE: 1 PYTHONUTF8: 1 - RUST_TOOLCHAIN: "1.97" jobs: lint: @@ -56,8 +55,8 @@ jobs: - name: Set up Rust run: | - rustup toolchain install "${RUST_TOOLCHAIN}" --profile minimal --component clippy --component rustfmt - rustup default "${RUST_TOOLCHAIN}" + rustup toolchain install stable --profile minimal --component clippy --component rustfmt + rustup default stable - name: Install dependencies run: uv sync --locked --group dev @@ -101,8 +100,8 @@ jobs: - name: Set up Rust run: | - rustup toolchain install "${RUST_TOOLCHAIN}" --profile minimal - rustup default "${RUST_TOOLCHAIN}" + rustup toolchain install stable --profile minimal + rustup default stable - name: Run tests for Python versions shell: bash @@ -174,8 +173,8 @@ jobs: - name: Set up Rust run: | - rustup toolchain install "${RUST_TOOLCHAIN}" --profile minimal - rustup default "${RUST_TOOLCHAIN}" + rustup toolchain install stable --profile minimal + rustup default stable - name: Verify tag matches package version if: github.event_name == 'push' && startsWith(github.ref, 'refs/tags/pynumaflow-lite-v') diff --git a/packages/pynumaflow-lite/manifests/batchmap/batchmap_cat.py b/packages/pynumaflow-lite/manifests/batchmap/batchmap_cat.py index 9dff944a..fed9eedf 100644 --- a/packages/pynumaflow-lite/manifests/batchmap/batchmap_cat.py +++ b/packages/pynumaflow-lite/manifests/batchmap/batchmap_cat.py @@ -1,20 +1,46 @@ -from collections.abc import AsyncIterator +import asyncio +import signal +from collections.abc import AsyncIterable, Awaitable, Callable from pynumaflow_lite import batchmapper -from pynumaflow_lite.batchmapper import BatchResponse, Datum, Message +from pynumaflow_lite.batchmapper import Message -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 - ] +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 if __name__ == "__main__": - batch_mapper_obj = SimpleBatchCat() - batchmapper.BatchMapAsyncServer(batch_mapper_obj.handler).run() + async_handler = SimpleBatchCat() + asyncio.run(start(async_handler)) diff --git a/packages/pynumaflow-lite/manifests/map/map_cat.py b/packages/pynumaflow-lite/manifests/map/map_cat.py index 8edc9cc5..f877f212 100644 --- a/packages/pynumaflow-lite/manifests/map/map_cat.py +++ b/packages/pynumaflow-lite/manifests/map/map_cat.py @@ -1,13 +1,48 @@ +import asyncio +import signal +from collections.abc import Awaitable, Callable + from pynumaflow_lite import mapper -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)] +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 if __name__ == "__main__": - mapper_obj = SimpleCat() - mapper.MapAsyncServer(mapper_obj.handler).run() + async_handler = SimpleCat() + asyncio.run(start(async_handler)) diff --git a/packages/pynumaflow-lite/manifests/mapstream/mapstream_cat.py b/packages/pynumaflow-lite/manifests/mapstream/mapstream_cat.py index 84933c98..9cfe0beb 100644 --- a/packages/pynumaflow-lite/manifests/mapstream/mapstream_cat.py +++ b/packages/pynumaflow-lite/manifests/mapstream/mapstream_cat.py @@ -1,19 +1,41 @@ -from collections.abc import AsyncIterable +import asyncio +import signal +from collections.abc import AsyncIterator, Callable from pynumaflow_lite import mapstreamer from pynumaflow_lite.mapstreamer import Message -class SimpleStreamCat: - async def handler(self, datum: mapstreamer.Datum) -> AsyncIterable[Message]: +class SimpleStreamCat(mapstreamer.MapStreamer): + async def handler(self, keys: list[str], datum: mapstreamer.Datum) -> AsyncIterator[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=datum.keys) + 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 if __name__ == "__main__": - map_streamer_obj = SimpleStreamCat() - mapstreamer.MapStreamAsyncServer(map_streamer_obj.handler).run() + async_handler = SimpleStreamCat() + asyncio.run(start(async_handler)) diff --git a/packages/pynumaflow-lite/manifests/sideinput/sideinput_example.py b/packages/pynumaflow-lite/manifests/sideinput/sideinput_example.py index 8fb790a9..5361b57f 100644 --- a/packages/pynumaflow-lite/manifests/sideinput/sideinput_example.py +++ b/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: +class SideInputHandler(mapper.Mapper): """ A Mapper that reads from side input files and includes the value in its output. """ @@ -51,11 +51,13 @@ class SideInputHandler: # Side input file that we are watching watched_file = "myticker" - async def handler(self, datum: mapper.Datum) -> list[mapper.Message]: + async def handler(self, keys: list[str], datum: mapper.Datum) -> mapper.Messages: with self.data_value_lock: current_value = self.data_value - return [mapper.Message(str.encode(current_value))] + messages = mapper.Messages() + messages.append(mapper.Message(str.encode(current_value))) + return messages def file_watcher(self): """ @@ -101,18 +103,27 @@ async def start_sideinput(): server.stop() -def start_mapper(): +async def start_mapper(): """Start the Mapper server that reads from side inputs.""" - mapper_obj = SideInputHandler() + server = mapper.MapAsyncServer() + handler = SideInputHandler() # Initialize the data value from the side input file - mapper_obj.init_data_value() + handler.init_data_value() # Start the file watcher in a background thread - watcher_thread = Thread(target=mapper_obj.file_watcher, daemon=True) + watcher_thread = Thread(target=handler.file_watcher, daemon=True) watcher_thread.start() - mapper.MapAsyncServer(mapper_obj.handler).run() + 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() if __name__ == "__main__": @@ -121,7 +132,7 @@ def start_mapper(): if is_mapper: print("Starting as Mapper (reading side inputs)...") - start_mapper() + asyncio.run(start_mapper()) else: print("Starting as SideInput retriever...") asyncio.run(start_sideinput()) diff --git a/packages/pynumaflow-lite/manifests/sink/sink_log.py b/packages/pynumaflow-lite/manifests/sink/sink_log.py index 1a94aeff..1b1c656c 100644 --- a/packages/pynumaflow-lite/manifests/sink/sink_log.py +++ b/packages/pynumaflow-lite/manifests/sink/sink_log.py @@ -1,19 +1,20 @@ import logging -from collections.abc import AsyncIterator +from collections.abc import AsyncIterable from pynumaflow_lite import sinker +from pynumaflow_lite.sinker import Sinker # Configure logging logging.basicConfig(level=logging.INFO) _LOGGER = logging.getLogger(__name__) -class SimpleLogSink: +class SimpleLogSink(Sinker): """ Simple log sink that logs each message and returns success responses. """ - async def handler(self, datums: AsyncIterator[sinker.Datum]) -> list[sinker.Response]: + async def handler(self, datums: AsyncIterable[sinker.Datum]) -> list[sinker.Response]: responses = [] async for msg in datums: _LOGGER.info("User Defined Sink: %s", msg.value.decode("utf-8")) @@ -24,5 +25,4 @@ async def handler(self, datums: AsyncIterator[sinker.Datum]) -> list[sinker.Resp if __name__ == "__main__": - sinker_obj = SimpleLogSink() - sinker.SinkAsyncServer(sinker_obj.handler).run() + sinker.SinkAsyncServer(SimpleLogSink()).run() diff --git a/packages/pynumaflow-lite/pynumaflow_lite/__init__.py b/packages/pynumaflow-lite/pynumaflow_lite/__init__.py index 30dab9f0..d1bfeb9f 100644 --- a/packages/pynumaflow-lite/pynumaflow_lite/__init__.py +++ b/packages/pynumaflow-lite/pynumaflow_lite/__init__.py @@ -61,27 +61,29 @@ except Exception: # pragma: no cover sideinputer = None -# Surface the Python async servers and data-type classes under the extension submodules for convenient access +# Surface the Python Mapper, BatchMapper, MapStreamer, Reducer, SessionReducer, ReduceStreamer, Accumulator, Sinker, +# Sourcer, SourceTransformer, and SideInput classes under the extension submodules for convenient access from ._accumulator_dtypes import Accumulator -from ._batchmap_server import BatchMapAsyncServer -from ._map_server import MapAsyncServer -from ._mapstream_server import MapStreamAsyncServer +from ._batchmapper_dtypes import BatchMapper +from ._map_dtypes import Mapper +from ._mapstream_dtypes import MapStreamer 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.MapAsyncServer = MapAsyncServer + mapper.Mapper = Mapper if batchmapper is not None: - batchmapper.BatchMapAsyncServer = BatchMapAsyncServer + batchmapper.BatchMapper = BatchMapper if mapstreamer is not None: - mapstreamer.MapStreamAsyncServer = MapStreamAsyncServer + mapstreamer.MapStreamer = MapStreamer if reducer is not None: reducer.Reducer = Reducer @@ -96,6 +98,7 @@ accumulator.Accumulator = Accumulator if sinker is not None: + sinker.Sinker = Sinker sinker.SinkAsyncServer = SinkAsyncServer if sourcer is not None: diff --git a/packages/pynumaflow-lite/pynumaflow_lite/_batchmap_server.py b/packages/pynumaflow-lite/pynumaflow_lite/_batchmap_server.py deleted file mode 100644 index b39f5164..00000000 --- a/packages/pynumaflow-lite/pynumaflow_lite/_batchmap_server.py +++ /dev/null @@ -1,91 +0,0 @@ -from __future__ import annotations - -import asyncio -import signal -from collections.abc import AsyncIterator, Awaitable, Callable -from types import TracebackType - -from .pynumaflow_lite import batchmapper as _batchmapper - -BatchResponse = _batchmapper.BatchResponse -Datum = _batchmapper.Datum - - -class BatchMapAsyncServer: - def __init__( - self, - handler: Callable[ - [AsyncIterator[Datum]], - Awaitable[list[BatchResponse]], - ], - *, - sock_file: str | None = None, - server_info_file: str | None = None, - ) -> None: - self._core = _batchmapper._BatchMapAsyncServer(sock_file, server_info_file) - self._handler = handler - self._task: asyncio.Task[None] | None = None - self._serving = False - - async def serve(self) -> None: - if self._serving: - raise RuntimeError("batchmap server is already serving") - self._serving = True - try: - await self._core.start(self._handler) - finally: - self._serving = False - - def stop(self) -> None: - self._core.stop() - - async def wait_ready(self, timeout: float = 30.0) -> None: - await self._core.wait_ready(timeout) - - async def __aenter__(self) -> BatchMapAsyncServer: - if self._task is not None and not self._task.done(): - raise RuntimeError("batchmap server is already serving") - - self._task = asyncio.create_task(self.serve()) - try: - await self.wait_ready() - except asyncio.CancelledError: - self.stop() - if self._task is not None: - await self._task - raise - except Exception: - self.stop() - if self._task is not None: - await self._task - raise - return self - - async def __aexit__( - self, - exc_type: type[BaseException] | None, - exc: BaseException | None, - tb: TracebackType | None, - ) -> None: - self.stop() - if self._task is not None: - try: - await self._task - finally: - self._task = None - - async def _main(self) -> None: - loop = asyncio.get_running_loop() - try: - loop.add_signal_handler(signal.SIGINT, self.stop) - loop.add_signal_handler(signal.SIGTERM, self.stop) - except (NotImplementedError, RuntimeError): - pass - - await self.serve() - - def run(self) -> None: - try: - asyncio.run(self._main()) - except KeyboardInterrupt: - self.stop() diff --git a/packages/pynumaflow-lite/pynumaflow_lite/_batchmapper_dtypes.py b/packages/pynumaflow-lite/pynumaflow_lite/_batchmapper_dtypes.py new file mode 100644 index 00000000..557c8afe --- /dev/null +++ b/packages/pynumaflow-lite/pynumaflow_lite/_batchmapper_dtypes.py @@ -0,0 +1,21 @@ +from abc import ABCMeta, abstractmethod +from collections.abc import AsyncIterable + +from pynumaflow_lite.batchmapper import BatchResponse, Datum + + +class BatchMapper(metaclass=ABCMeta): + """ + Provides an interface to write a BatchMap servicer. + """ + + def __call__(self, *args, **kwargs): + return self.handler(*args, **kwargs) + + @abstractmethod + async def handler(self, batch: AsyncIterable[Datum]) -> list[BatchResponse]: + """ + Implement this handler function for batch mapping. + The returned list length should equal the input batch size. + """ + pass diff --git a/packages/pynumaflow-lite/pynumaflow_lite/_map_dtypes.py b/packages/pynumaflow-lite/pynumaflow_lite/_map_dtypes.py new file mode 100644 index 00000000..13733ed7 --- /dev/null +++ b/packages/pynumaflow-lite/pynumaflow_lite/_map_dtypes.py @@ -0,0 +1,24 @@ +from abc import ABCMeta, abstractmethod + +from pynumaflow_lite.mapper import Datum, Messages + + +class Mapper(metaclass=ABCMeta): + """ + Provides an interface to write a SyncMapServicer + which will be exposed over a Synchronous gRPC server. + """ + + def __call__(self, *args, **kwargs): + """ + This allows to execute the handler function directly if + class instance is sent as a callable. + """ + return self.handler(*args, **kwargs) + + @abstractmethod + async def handler(self, keys: list[str], payload: Datum) -> Messages: + """ + Implement this handler function which implements the MapAsyncCallable interface. + """ + pass diff --git a/packages/pynumaflow-lite/pynumaflow_lite/_map_server.py b/packages/pynumaflow-lite/pynumaflow_lite/_map_server.py deleted file mode 100644 index 4b3cb68e..00000000 --- a/packages/pynumaflow-lite/pynumaflow_lite/_map_server.py +++ /dev/null @@ -1,88 +0,0 @@ -from __future__ import annotations - -import asyncio -import signal -from collections.abc import Awaitable, Callable -from types import TracebackType - -from .pynumaflow_lite import mapper as _mapper - -Datum = _mapper.Datum -Message = _mapper.Message - - -class MapAsyncServer: - def __init__( - self, - handler: Callable[[Datum], Awaitable[list[Message]]], - *, - sock_file: str | None = None, - server_info_file: str | None = None, - ) -> None: - self._core = _mapper._MapAsyncServer(sock_file, server_info_file) - self._handler = handler - self._task: asyncio.Task[None] | None = None - self._serving = False - - async def serve(self) -> None: - if self._serving: - raise RuntimeError("map server is already serving") - self._serving = True - try: - await self._core.start(self._handler) - finally: - self._serving = False - - def stop(self) -> None: - self._core.stop() - - async def wait_ready(self, timeout: float = 30.0) -> None: - await self._core.wait_ready(timeout) - - async def __aenter__(self) -> MapAsyncServer: - if self._task is not None and not self._task.done(): - raise RuntimeError("map server is already serving") - - self._task = asyncio.create_task(self.serve()) - try: - await self.wait_ready() - except asyncio.CancelledError: - self.stop() - if self._task is not None: - await self._task - raise - except Exception: - self.stop() - if self._task is not None: - await self._task - raise - return self - - async def __aexit__( - self, - exc_type: type[BaseException] | None, - exc: BaseException | None, - tb: TracebackType | None, - ) -> None: - self.stop() - if self._task is not None: - try: - await self._task - finally: - self._task = None - - async def _main(self) -> None: - loop = asyncio.get_running_loop() - try: - loop.add_signal_handler(signal.SIGINT, self.stop) - loop.add_signal_handler(signal.SIGTERM, self.stop) - except (NotImplementedError, RuntimeError): - pass - - await self.serve() - - def run(self) -> None: - try: - asyncio.run(self._main()) - except KeyboardInterrupt: - self.stop() diff --git a/packages/pynumaflow-lite/pynumaflow_lite/_mapstream_dtypes.py b/packages/pynumaflow-lite/pynumaflow_lite/_mapstream_dtypes.py new file mode 100644 index 00000000..32c2b9cb --- /dev/null +++ b/packages/pynumaflow-lite/pynumaflow_lite/_mapstream_dtypes.py @@ -0,0 +1,22 @@ +from abc import ABCMeta, abstractmethod +from collections.abc import AsyncIterator + +from pynumaflow_lite.mapstreamer import Datum, Message + + +class MapStreamer(metaclass=ABCMeta): + """ + Provides an interface to write a streaming map servicer. + The handler yields outputs incrementally as an async iterator. + """ + + def __call__(self, *args, **kwargs): + return self.handler(*args, **kwargs) + + @abstractmethod + async def handler(self, keys: list[str], datum: Datum) -> AsyncIterator[Message]: + """ + Implement this handler function for streaming mapping. + It should be an async generator yielding Message objects. + """ + pass diff --git a/packages/pynumaflow-lite/pynumaflow_lite/_mapstream_server.py b/packages/pynumaflow-lite/pynumaflow_lite/_mapstream_server.py deleted file mode 100644 index 4349878d..00000000 --- a/packages/pynumaflow-lite/pynumaflow_lite/_mapstream_server.py +++ /dev/null @@ -1,88 +0,0 @@ -from __future__ import annotations - -import asyncio -import signal -from collections.abc import AsyncIterable, Callable -from types import TracebackType - -from .pynumaflow_lite import mapstreamer as _mapstreamer - -Datum = _mapstreamer.Datum -Message = _mapstreamer.Message - - -class MapStreamAsyncServer: - def __init__( - self, - handler: Callable[[Datum], AsyncIterable[Message]], - *, - sock_file: str | None = None, - server_info_file: str | None = None, - ) -> None: - self._core = _mapstreamer._MapStreamAsyncServer(sock_file, server_info_file) - self._handler = handler - self._task: asyncio.Task[None] | None = None - self._serving = False - - async def serve(self) -> None: - if self._serving: - raise RuntimeError("mapstream server is already serving") - self._serving = True - try: - await self._core.start(self._handler) - finally: - self._serving = False - - def stop(self) -> None: - self._core.stop() - - async def wait_ready(self, timeout: float = 30.0) -> None: - await self._core.wait_ready(timeout) - - async def __aenter__(self) -> MapStreamAsyncServer: - if self._task is not None and not self._task.done(): - raise RuntimeError("mapstream server is already serving") - - self._task = asyncio.create_task(self.serve()) - try: - await self.wait_ready() - except asyncio.CancelledError: - self.stop() - if self._task is not None: - await self._task - raise - except Exception: - self.stop() - if self._task is not None: - await self._task - raise - return self - - async def __aexit__( - self, - exc_type: type[BaseException] | None, - exc: BaseException | None, - tb: TracebackType | None, - ) -> None: - self.stop() - if self._task is not None: - try: - await self._task - finally: - self._task = None - - async def _main(self) -> None: - loop = asyncio.get_running_loop() - try: - loop.add_signal_handler(signal.SIGINT, self.stop) - loop.add_signal_handler(signal.SIGTERM, self.stop) - except (NotImplementedError, RuntimeError): - pass - - await self.serve() - - def run(self) -> None: - try: - asyncio.run(self._main()) - except KeyboardInterrupt: - self.stop() diff --git a/packages/pynumaflow-lite/pynumaflow_lite/_sink_dtypes.py b/packages/pynumaflow-lite/pynumaflow_lite/_sink_dtypes.py new file mode 100644 index 00000000..840dd65f --- /dev/null +++ b/packages/pynumaflow-lite/pynumaflow_lite/_sink_dtypes.py @@ -0,0 +1,21 @@ +from abc import ABCMeta, abstractmethod +from collections.abc import AsyncIterable + +from pynumaflow_lite.sinker import Datum, Response + + +class Sinker(metaclass=ABCMeta): + """ + Provides an interface to write a Sink servicer. + """ + + def __call__(self, *args, **kwargs): + return self.handler(*args, **kwargs) + + @abstractmethod + async def handler(self, datums: AsyncIterable[Datum]) -> list[Response]: + """ + Implement this handler function for sink. + Process the stream of datums and return responses. + """ + pass diff --git a/packages/pynumaflow-lite/pynumaflow_lite/_sink_server.py b/packages/pynumaflow-lite/pynumaflow_lite/_sink_server.py index 5a4ef95d..4dd29892 100644 --- a/packages/pynumaflow-lite/pynumaflow_lite/_sink_server.py +++ b/packages/pynumaflow-lite/pynumaflow_lite/_sink_server.py @@ -2,22 +2,16 @@ import asyncio import signal -from collections.abc import AsyncIterator, Awaitable, Callable from types import TracebackType +from typing import Any from .pynumaflow_lite import sinker as _sinker -Datum = _sinker.Datum -Response = _sinker.Response - class SinkAsyncServer: def __init__( self, - handler: Callable[ - [AsyncIterator[Datum]], - Awaitable[list[Response]], - ], + handler: Any, *, sock_file: str | None = None, server_info_file: str | None = None, @@ -25,16 +19,9 @@ def __init__( self._core = _sinker._SinkAsyncServer(sock_file, server_info_file) self._handler = handler self._task: asyncio.Task[None] | None = None - self._serving = False async def serve(self) -> None: - if self._serving: - raise RuntimeError("sink server is already serving") - self._serving = True - try: - await self._core.start(self._handler) - finally: - self._serving = False + await self._core.start(self._handler) def stop(self) -> None: self._core.stop() diff --git a/packages/pynumaflow-lite/pynumaflow_lite/batchmapper.pyi b/packages/pynumaflow-lite/pynumaflow_lite/batchmapper.pyi index 71cbcf9c..277af5e3 100644 --- a/packages/pynumaflow-lite/pynumaflow_lite/batchmapper.pyi +++ b/packages/pynumaflow-lite/pynumaflow_lite/batchmapper.pyi @@ -2,7 +2,6 @@ from __future__ import annotations import datetime as _dt from collections.abc import AsyncIterator, Awaitable, Callable -from types import TracebackType class Message: keys: list[str] | None @@ -16,77 +15,48 @@ class Message: tags: list[str] | None = ..., ) -> None: ... @staticmethod - def to_drop() -> Message: ... - def __repr__(self) -> str: ... - def __eq__(self, other: object) -> bool: ... + def message_to_drop() -> Message: ... class Datum: keys: list[str] value: bytes watermark: _dt.datetime - event_time: _dt.datetime + eventtime: _dt.datetime id: str headers: dict[str, str] - def __init__( - self, - *, - keys: list[str] = ..., - value: bytes = ..., - id: str = ..., - event_time: _dt.datetime | None = ..., - watermark: _dt.datetime | None = ..., - headers: dict[str, str] = ..., - ) -> None: ... def __repr__(self) -> str: ... def __str__(self) -> str: ... class BatchResponse: id: str - messages: list[Message] - def __init__(self, id: str, messages: list[Message] | None = ...) -> None: ... + def __init__(self, id: str) -> None: ... + @staticmethod + def from_id(id: str) -> BatchResponse: ... def append(self, message: Message) -> None: ... - def __len__(self) -> int: ... - def __repr__(self) -> str: ... - def __eq__(self, other: object) -> bool: ... -class _BatchMapAsyncServer: - def __init__( - self, - sock_file: str | None = ..., - server_info_file: str | None = ..., - ) -> None: ... - def start( - self, - handler: Callable[[AsyncIterator[Datum]], Awaitable[list[BatchResponse]]], - ) -> Awaitable[None]: ... - def wait_ready(self, timeout: float = ...) -> Awaitable[None]: ... - def stop(self) -> None: ... +class BatchResponses: + def __init__(self) -> None: ... + def append(self, response: BatchResponse) -> None: ... class BatchMapAsyncServer: def __init__( self, - handler: Callable[[AsyncIterator[Datum]], Awaitable[list[BatchResponse]]], - *, sock_file: str | None = ..., - server_info_file: str | None = ..., + info_file: str | None = ..., ) -> None: ... - def run(self) -> None: ... - async def serve(self) -> None: ... + def start(self, py_func: Callable[[AsyncIterator[Datum]], Awaitable[BatchResponses]]) -> Awaitable[None]: ... def stop(self) -> None: ... - async def wait_ready(self, timeout: float = ...) -> None: ... - async def __aenter__(self) -> BatchMapAsyncServer: ... - async def __aexit__( - self, - exc_type: type[BaseException] | None, - exc: BaseException | None, - tb: TracebackType | None, - ) -> None: ... + +class BatchMapper: + async def handler(self, batch: AsyncIterator[Datum]) -> BatchResponses: ... __all__ = [ "BatchMapAsyncServer", + "BatchMapper", "BatchResponse", + "BatchResponses", "Datum", "Message", ] diff --git a/packages/pynumaflow-lite/pynumaflow_lite/mapper.pyi b/packages/pynumaflow-lite/pynumaflow_lite/mapper.pyi index b5c432a6..8255e6d6 100644 --- a/packages/pynumaflow-lite/pynumaflow_lite/mapper.pyi +++ b/packages/pynumaflow-lite/pynumaflow_lite/mapper.pyi @@ -2,84 +2,114 @@ from __future__ import annotations import datetime as _dt from collections.abc import Awaitable, Callable -from types import TracebackType + +# Re-export the Python ABC for user convenience and typing +from ._map_dtypes import Mapper as Mapper + +class SystemMetadata: + """System-generated metadata groups per message (read-only).""" + + def __init__(self) -> None: ... + def groups(self) -> list[str]: + """Returns the groups of the system metadata.""" + ... + + def keys(self, group: str) -> list[str]: + """Returns the keys of the system metadata for the given group.""" + ... + + def value(self, group: str, key: str) -> bytes: + """Returns the value of the system metadata for the given group and key.""" + ... + + def __repr__(self) -> str: ... + +class UserMetadata: + """User-defined metadata groups per message (read-write).""" + + def __init__(self) -> None: ... + def groups(self) -> list[str]: + """Returns the groups of the user metadata.""" + ... + + def keys(self, group: str) -> list[str]: + """Returns the keys of the user metadata for the given group.""" + ... + + def value(self, group: str, key: str) -> bytes: + """Returns the value of the user metadata for the given group and key.""" + ... + + def create_group(self, group: str) -> None: + """Creates a new group in the user metadata.""" + ... + + def add_kv(self, group: str, key: str, value: bytes) -> None: + """Adds a key-value pair to the user metadata.""" + ... + + def remove_key(self, group: str, key: str) -> None: + """Removes a key from a group in the user metadata.""" + ... + + def remove_group(self, group: str) -> None: + """Removes a group from the user metadata.""" + ... + + def __repr__(self) -> str: ... + +class Messages: + def __init__(self) -> None: ... + def append(self, message: Message) -> None: ... + def __repr__(self) -> str: ... + def __str__(self) -> str: ... class Message: keys: list[str] | None value: bytes tags: list[str] | None - user_metadata: dict[str, dict[str, bytes]] | None + user_metadata: UserMetadata | None def __init__( self, value: bytes, keys: list[str] | None = ..., tags: list[str] | None = ..., - user_metadata: dict[str, dict[str, bytes]] | None = ..., + user_metadata: UserMetadata | None = ..., ) -> None: ... @staticmethod - def to_drop() -> Message: ... - def __repr__(self) -> str: ... - def __eq__(self, other: object) -> bool: ... + def message_to_drop() -> Message: ... class Datum: + # Read-only attributes provided by the extension keys: list[str] value: bytes watermark: _dt.datetime - event_time: _dt.datetime + eventtime: _dt.datetime headers: dict[str, str] - user_metadata: dict[str, dict[str, bytes]] - system_metadata: dict[str, dict[str, bytes]] + user_metadata: UserMetadata + system_metadata: SystemMetadata - def __init__( - self, - *, - keys: list[str] = ..., - value: bytes = ..., - event_time: _dt.datetime | None = ..., - watermark: _dt.datetime | None = ..., - headers: dict[str, str] = ..., - user_metadata: dict[str, dict[str, bytes]] = ..., - system_metadata: dict[str, dict[str, bytes]] = ..., - ) -> None: ... def __repr__(self) -> str: ... def __str__(self) -> str: ... -class _MapAsyncServer: - def __init__( - self, - sock_file: str | None = ..., - server_info_file: str | None = ..., - ) -> None: ... - def start( - self, - handler: Callable[[Datum], Awaitable[list[Message]]], - ) -> Awaitable[None]: ... - def wait_ready(self, timeout: float = ...) -> Awaitable[None]: ... - def stop(self) -> None: ... - class MapAsyncServer: def __init__( self, - handler: Callable[[Datum], Awaitable[list[Message]]], - *, sock_file: str | None = ..., - server_info_file: str | None = ..., + info_file: str | None = ..., ) -> None: ... - def run(self) -> None: ... - async def serve(self) -> None: ... + def start(self, py_func: Callable[[list[str], Datum], Awaitable[Messages]]) -> Awaitable[None]: ... def stop(self) -> None: ... - async def wait_ready(self, timeout: float = ...) -> None: ... - async def __aenter__(self) -> MapAsyncServer: ... - async def __aexit__( - self, - exc_type: type[BaseException] | None, - exc: BaseException | None, - tb: TracebackType | None, - ) -> None: ... + +# Simple utility function exposed by the extension __all__ = [ "Datum", "MapAsyncServer", + "Mapper", "Message", + "Messages", + "SystemMetadata", + "UserMetadata", ] diff --git a/packages/pynumaflow-lite/pynumaflow_lite/mapstreamer.pyi b/packages/pynumaflow-lite/pynumaflow_lite/mapstreamer.pyi index d6660aa3..ec9209af 100644 --- a/packages/pynumaflow-lite/pynumaflow_lite/mapstreamer.pyi +++ b/packages/pynumaflow-lite/pynumaflow_lite/mapstreamer.pyi @@ -1,8 +1,7 @@ from __future__ import annotations import datetime as _dt -from collections.abc import AsyncIterable, Awaitable, Callable -from types import TracebackType +from collections.abc import AsyncIterator, Awaitable, Callable class Message: keys: list[str] | None @@ -16,64 +15,35 @@ class Message: tags: list[str] | None = ..., ) -> None: ... @staticmethod + def message_to_drop() -> Message: ... + @staticmethod def to_drop() -> Message: ... - def __repr__(self) -> str: ... - def __eq__(self, other: object) -> bool: ... class Datum: keys: list[str] value: bytes watermark: _dt.datetime - event_time: _dt.datetime + eventtime: _dt.datetime headers: dict[str, str] - def __init__( - self, - *, - keys: list[str] = ..., - value: bytes = ..., - event_time: _dt.datetime | None = ..., - watermark: _dt.datetime | None = ..., - headers: dict[str, str] = ..., - ) -> None: ... def __repr__(self) -> str: ... def __str__(self) -> str: ... -class _MapStreamAsyncServer: - def __init__( - self, - sock_file: str | None = ..., - server_info_file: str | None = ..., - ) -> None: ... - def start( - self, - handler: Callable[[Datum], AsyncIterable[Message]], - ) -> Awaitable[None]: ... - def wait_ready(self, timeout: float = ...) -> Awaitable[None]: ... - def stop(self) -> None: ... - class MapStreamAsyncServer: def __init__( self, - handler: Callable[[Datum], AsyncIterable[Message]], - *, sock_file: str | None = ..., - server_info_file: str | None = ..., + info_file: str | None = ..., ) -> None: ... - def run(self) -> None: ... - async def serve(self) -> None: ... + def start(self, py_func: Callable[[list[str], Datum], AsyncIterator[Message]]) -> Awaitable[None]: ... def stop(self) -> None: ... - async def wait_ready(self, timeout: float = ...) -> None: ... - async def __aenter__(self) -> MapStreamAsyncServer: ... - async def __aexit__( - self, - exc_type: type[BaseException] | None, - exc: BaseException | None, - tb: TracebackType | None, - ) -> None: ... + +class MapStreamer: + async def handler(self, keys: list[str], datum: Datum) -> AsyncIterator[Message]: ... __all__ = [ "Datum", "MapStreamAsyncServer", + "MapStreamer", "Message", ] diff --git a/packages/pynumaflow-lite/pynumaflow_lite/sinker.pyi b/packages/pynumaflow-lite/pynumaflow_lite/sinker.pyi index 5be350e0..5278f341 100644 --- a/packages/pynumaflow-lite/pynumaflow_lite/sinker.pyi +++ b/packages/pynumaflow-lite/pynumaflow_lite/sinker.pyi @@ -60,23 +60,26 @@ class Datum: def __repr__(self) -> str: ... def __str__(self) -> str: ... +_SinkHandler = Callable[[AsyncIterator[Datum]], Awaitable[list[Response]]] + class _SinkAsyncServer: def __init__( self, sock_file: str | None = ..., server_info_file: str | None = ..., ) -> None: ... - def start( - self, - handler: Callable[[AsyncIterator[Datum]], Awaitable[list[Response]]], - ) -> Awaitable[None]: ... + def start(self, handler: _SinkHandler) -> Awaitable[None]: ... def wait_ready(self, timeout: float = ...) -> Awaitable[None]: ... def stop(self) -> None: ... +class Sinker: + def __call__(self, datums: AsyncIterator[Datum]) -> Awaitable[list[Response]]: ... + async def handler(self, datums: AsyncIterator[Datum]) -> list[Response]: ... + class SinkAsyncServer: def __init__( self, - handler: Callable[[AsyncIterator[Datum]], Awaitable[list[Response]]], + handler: _SinkHandler | Sinker, *, sock_file: str | None = ..., server_info_file: str | None = ..., @@ -98,4 +101,5 @@ __all__ = [ "Message", "Response", "SinkAsyncServer", + "Sinker", ] diff --git a/packages/pynumaflow-lite/src/batchmap/mod.rs b/packages/pynumaflow-lite/src/batchmap/mod.rs index e127110a..fc43433b 100644 --- a/packages/pynumaflow-lite/src/batchmap/mod.rs +++ b/packages/pynumaflow-lite/src/batchmap/mod.rs @@ -1,5 +1,4 @@ use std::collections::HashMap; -use std::time::Duration; use numaflow::batchmap; @@ -15,18 +14,15 @@ use pyo3::prelude::*; use std::sync::Mutex; /// A message to be sent for a single datum in batch response. -#[pyclass(module = "pynumaflow_lite.batchmapper", from_py_object, eq)] -#[derive(Clone, Default, Debug, PartialEq)] +#[pyclass(module = "pynumaflow_lite.batchmapper", from_py_object)] +#[derive(Clone, Default, Debug)] pub struct Message { /// Keys are a collection of strings which will be passed on to the next vertex as is. It can /// be an empty collection. - #[pyo3(get)] pub keys: Option>, /// Value is the value passed to the next vertex. - #[pyo3(get)] pub value: Vec, /// Tags are used for conditional forwarding. - #[pyo3(get)] pub tags: Option>, } @@ -41,28 +37,15 @@ impl Message { } /// Drop a [Message], do not forward to the next vertex. + #[pyo3(signature = ())] #[staticmethod] - #[pyo3(signature = () -> "Message")] - fn to_drop() -> Self { + fn message_to_drop() -> Self { Self { keys: None, value: vec![], tags: Some(vec![numaflow::shared::DROP.to_string()]), } } - - fn __repr__(&self) -> String { - format!( - "Message(value={}, keys={}, tags={})", - crate::map::bytes_literal(&self.value), - self.keys - .as_ref() - .map_or_else(|| "None".to_string(), |keys| format!("{keys:?}")), - self.tags - .as_ref() - .map_or_else(|| "None".to_string(), |tags| format!("{tags:?}")), - ) - } } impl From for batchmap::Message { @@ -89,7 +72,7 @@ pub struct Datum { pub watermark: DateTime, /// Time of the element as seen at source or aligned after a reduce operation. #[pyo3(get)] - pub event_time: DateTime, + pub eventtime: DateTime, /// ID is the unique id of the message #[pyo3(get)] pub id: String, @@ -98,104 +81,112 @@ pub struct Datum { pub headers: HashMap, } -#[pymethods] impl Datum { - #[new] - #[pyo3(signature = ( - *, - keys: "list[str] | None"=None, - value: "bytes | None"=None, - id: "str | None"=None, - event_time: "datetime.datetime | None"=None, - watermark: "datetime.datetime | None"=None, - headers: "dict[str, str] | None"=None, - ) -> "Datum")] fn new( - keys: Option>, - value: Option>, - id: Option, - event_time: Option>, - watermark: Option>, - headers: Option>, + keys: Vec, + value: Vec, + watermark: DateTime, + eventtime: DateTime, + id: String, + headers: HashMap, ) -> Self { Self { - keys: keys.unwrap_or_default(), - value: value.unwrap_or_default(), - watermark: watermark.unwrap_or(DateTime::::UNIX_EPOCH), - event_time: event_time.unwrap_or(DateTime::::UNIX_EPOCH), - id: id.unwrap_or_default(), - headers: headers.unwrap_or_default(), + keys, + value, + watermark, + eventtime, + id, + headers, } } fn __repr__(&self) -> String { format!( - "Datum(keys={:?}, value={}, watermark={}, event_time={}, id={:?}, headers={:?})", + "Datum(keys={:?}, value={:?}, watermark={}, eventtime={}, id={}, headers={:?})", + self.keys, self.value, self.watermark, self.eventtime, self.id, self.headers + ) + } + + fn __str__(&self) -> String { + format!( + "Datum(keys={:?}, value={:?}, watermark={}, eventtime={}, id={}, headers={:?})", self.keys, - crate::map::bytes_literal(&self.value), + String::from_utf8_lossy(&self.value), self.watermark, - self.event_time, + self.eventtime, self.id, self.headers ) } - - fn __str__(&self) -> String { - self.__repr__() - } } impl From for Datum { fn from(value: batchmap::Datum) -> Self { - Self { - keys: value.keys, - value: value.value, - watermark: value.watermark, - event_time: value.event_time, - id: value.id, - headers: value.headers, - } + Datum::new( + value.keys, + value.value, + value.watermark, + value.event_time, + value.id, + value.headers, + ) } } /// BatchResponse mirrors numaflow::batchmap::BatchResponse for Python -#[pyclass(module = "pynumaflow_lite.batchmapper", from_py_object, eq)] -#[derive(Clone, Debug, PartialEq)] +#[pyclass(module = "pynumaflow_lite.batchmapper", from_py_object)] +#[derive(Clone, Debug)] pub struct BatchResponse { #[pyo3(get)] pub id: String, - #[pyo3(get)] pub messages: Vec, } #[pymethods] impl BatchResponse { #[new] - #[pyo3(signature = (id: "str", messages: "list[Message] | None"=None) -> "BatchResponse")] - fn new(id: String, messages: Option>) -> Self { + #[pyo3(signature = (id: "str") -> "BatchResponse")] + fn new(id: String) -> Self { + Self { + id, + messages: Vec::new(), + } + } + + #[staticmethod] + #[pyo3(signature = (id: "str") -> "BatchResponse")] + fn from_id(id: String) -> Self { Self { id, - messages: messages.unwrap_or_default(), + messages: Vec::new(), } } - #[pyo3(signature = (message: "Message"))] + #[pyo3(signature = (message))] fn append(&mut self, message: Message) { self.messages.push(message); } +} + +/// A collection of BatchResponse objects for a batch. +#[pyclass(module = "pynumaflow_lite.batchmapper", from_py_object)] +#[derive(Clone, Debug)] +pub struct BatchResponses { + pub(crate) responses: Vec, +} - fn __len__(&self) -> usize { - self.messages.len() +#[pymethods] +impl BatchResponses { + #[new] + #[pyo3(signature = () -> "BatchResponses")] + fn new() -> Self { + Self { responses: vec![] } } - fn __repr__(&self) -> String { - let messages = self - .messages - .iter() - .map(|message| message.__repr__()) - .collect::>() - .join(", "); - format!("BatchResponse(id={:?}, messages=[{}])", self.id, messages) + /// Append a BatchResponse to the collection. + #[pyo3(signature = (response: "BatchResponse"))] + fn append(&mut self, response: BatchResponse) { + self.responses.push(response); } } @@ -218,6 +209,14 @@ pub struct PyAsyncDatumStream { #[pymethods] impl PyAsyncDatumStream { + #[new] + fn new() -> Self { + let (_tx, rx) = mpsc::channel::(1); + Self { + inner: crate::pyiterables::AsyncChannelStream::new(rx), + } + } + fn __aiter__(slf: PyRef<'_, Self>) -> PyRef<'_, Self> { slf } @@ -236,33 +235,30 @@ impl PyAsyncDatumStream { } /// Async Batch Map Server that can be started from Python code -#[pyclass(name = "_BatchMapAsyncServer", module = "pynumaflow_lite.batchmapper")] +#[pyclass(module = "pynumaflow_lite.batchmapper")] pub struct BatchMapAsyncServer { sock_file: String, - server_info_file: String, + info_file: String, shutdown_tx: Mutex>>, } #[pymethods] impl BatchMapAsyncServer { #[new] - #[pyo3(signature = ( - sock_file: "str | None"=None, - server_info_file: "str | None"=None, - ) -> "_BatchMapAsyncServer")] - fn new(sock_file: Option, server_info_file: Option) -> Self { + #[pyo3(signature = (sock_file: "str | None"=batchmap::SOCK_ADDR.to_string(), info_file: "str | None"=batchmap::SERVER_INFO_FILE.to_string()) -> "BatchMapAsyncServer" + )] + fn new(sock_file: String, info_file: String) -> Self { Self { - sock_file: sock_file.unwrap_or_else(|| batchmap::SOCK_ADDR.to_string()), - server_info_file: server_info_file - .unwrap_or_else(|| batchmap::SERVER_INFO_FILE.to_string()), + sock_file, + info_file, shutdown_tx: Mutex::new(None), } } - #[pyo3(signature = (handler: "callable") -> "None")] - pub fn start<'a>(&self, py: Python<'a>, handler: Py) -> PyResult> { + #[pyo3(signature = (py_func: "callable") -> "None")] + pub fn start<'a>(&self, py: Python<'a>, py_func: Py) -> PyResult> { let sock_file = self.sock_file.clone(); - let server_info_file = self.server_info_file.clone(); + let info_file = self.info_file.clone(); let (tx, rx) = tokio::sync::oneshot::channel::<()>(); { let mut guard = self.shutdown_tx.lock().unwrap(); @@ -270,25 +266,11 @@ impl BatchMapAsyncServer { } pyo3_async_runtimes::tokio::future_into_py(py, async move { - crate::batchmap::server::start(handler, sock_file, server_info_file, rx).await?; - Ok(()) - }) - } - - /// Wait until the Numaflow IsReady probe succeeds over the batchmap UDS. - #[pyo3(signature = (timeout: "float"=30.0) -> "None")] - pub fn wait_ready<'a>(&self, py: Python<'a>, timeout: f64) -> PyResult> { - if !timeout.is_finite() || timeout < 0.0 { - return Err(pyo3::PyErr::new::( - "timeout must be a non-negative finite float", - )); - } - - let sock_file = self.sock_file.clone(); - let timeout = Duration::from_secs_f64(timeout); - - pyo3_async_runtimes::tokio::future_into_py(py, async move { - crate::map::wait_for_ready(sock_file, timeout, "batchmap").await?; + // batch server uses the same runner loop and shutdown composition for now + // dedicated start is wired below + crate::batchmap::server::start(py_func, sock_file, info_file, rx) + .await + .expect("server failed to start"); Ok(()) }) } @@ -307,7 +289,9 @@ pub(crate) fn populate_py_module(m: &Bound) -> PyResult<()> { m.add_class::()?; m.add_class::()?; m.add_class::()?; + m.add_class::()?; m.add_class::()?; + m.add_class::()?; Ok(()) } diff --git a/packages/pynumaflow-lite/src/batchmap/server.rs b/packages/pynumaflow-lite/src/batchmap/server.rs index dd799146..25c75210 100644 --- a/packages/pynumaflow-lite/src/batchmap/server.rs +++ b/packages/pynumaflow-lite/src/batchmap/server.rs @@ -1,34 +1,13 @@ +// use crate::batchmap::Datum; use numaflow::batchmap; use numaflow::shared::ServerExtras; -use pyo3::exceptions::PyTypeError; use pyo3::prelude::*; -use std::sync::{Arc, Mutex}; +use std::sync::Arc; pub(crate) struct PyBatchMapRunner { pub(crate) event_loop: Arc>, pub(crate) py_func: Arc>, - pub(crate) error_slot: Arc>>, - pub(crate) shutdown_tx: Arc>>>, -} - -impl PyBatchMapRunner { - fn fail(&self, error: PyErr) -> Vec { - // Only the first error is reported; later batches may still be in flight - // while shutdown is underway, and their failures would be duplicates. - let mut error_slot = self.error_slot.lock().unwrap(); - if error_slot.is_none() { - Python::attach(|py| error.print(py)); - *error_slot = Some(error); - } - drop(error_slot); - - if let Some(tx) = self.shutdown_tx.lock().unwrap().take() { - let _ = tx.send(()); - } - - Vec::new() - } } #[tonic::async_trait] @@ -50,58 +29,28 @@ impl batchmap::BatchMapper for PyBatchMapRunner { // When input ends, dropping tx closes the channel }); - // Call the Python coroutine: py_func(batch: AsyncIterator[Datum]) -> list[BatchResponse] - let fut = match Python::attach(|py| -> PyResult<_> { + // Call the Python coroutine: py_func(batch: AsyncIterable[Datum]) -> BatchResponses + let fut = Python::attach(|py| { let locals = pyo3_async_runtimes::TaskLocals::new(self.event_loop.bind(py).clone()); let py_func = self.py_func.clone(); let stream = crate::batchmap::PyAsyncDatumStream::new_with(rx); - let coro = py_func.call1(py, (stream,))?.into_bound(py); - let is_awaitable: bool = py - .import("inspect")? - .call_method1("isawaitable", (&coro,))? - .extract()?; - if !is_awaitable { - return Err(PyErr::new::( - "batchmap handler must be an async function (coroutine)", - )); - } - pyo3_async_runtimes::into_future_with_locals(&locals, coro).map_err(|_| { - PyErr::new::( - "batchmap handler must be an async function (coroutine)", - ) - }) - }) { - Ok(fut) => fut, - Err(error) => return self.fail(error), - }; - - let result = match fut.await { - Ok(result) => result, - Err(error) => return self.fail(error), - }; + let coro = py_func.call1(py, (stream,)).unwrap().into_bound(py); + pyo3_async_runtimes::into_future_with_locals(&locals, coro).unwrap() + }); + + let result = fut.await.unwrap(); // Ensure forwarder completes let _ = forwarder.await; - let responses: Vec = match Python::attach(|py| { - result.extract(py).map_err(|_| { - let type_name = result - .bind(py) - .get_type() - .name() - .map(|name| name.to_string_lossy().into_owned()) - .unwrap_or_else(|_| "".to_string()); - PyErr::new::(format!( - "batchmap handler must return list[BatchResponse], got {type_name}" - )) - }) - }) { - Ok(responses) => responses, - Err(error) => return self.fail(error), - }; + let responses = Python::attach(|py| { + let x: crate::batchmap::BatchResponses = result.extract(py).unwrap(); + x + }); responses + .responses .into_iter() .map(|resp| resp.into()) .collect::>() @@ -116,34 +65,14 @@ pub(super) async fn start( shutdown_rx: tokio::sync::oneshot::Receiver<()>, ) -> Result<(), pyo3::PyErr> { let (tx, rx) = tokio::sync::oneshot::channel(); - let py_asyncio_loop_handle = tokio::task::spawn_blocking({ - println!( - "Starting BatchMap UDF. socket={}, server_info={}", - sock_file, info_file - ); - move || crate::pyrs::run_asyncio(tx) - }); + let py_asyncio_loop_handle = tokio::task::spawn_blocking(move || crate::pyrs::run_asyncio(tx)); let event_loop = rx.await.unwrap(); - let error_slot = Arc::new(Mutex::new(None)); - let (internal_shutdown_tx, internal_shutdown_rx) = tokio::sync::oneshot::channel(); - let (server_shutdown_tx, server_shutdown_rx) = tokio::sync::oneshot::channel(); - - tokio::spawn(async move { - tokio::select! { - _ = shutdown_rx => {}, - _ = internal_shutdown_rx => {}, - } - let _ = server_shutdown_tx.send(()); - }); - - let (sig_handle, combined_rx) = crate::pyrs::setup_sig_handler(server_shutdown_rx); + let (sig_handle, combined_rx) = crate::pyrs::setup_sig_handler(shutdown_rx); let py_runner = PyBatchMapRunner { py_func: Arc::new(py_func), event_loop: event_loop.clone(), - error_slot: error_slot.clone(), - shutdown_tx: Arc::new(Mutex::new(Some(internal_shutdown_tx))), }; let server = numaflow::batchmap::Server::new(py_runner) @@ -162,7 +91,7 @@ pub(super) async fn start( } }); - println!("Numaflow BatchMap has shutdown..."); + println!("Numaflow Core (batch) has shutdown..."); // Wait for the blocking asyncio thread to finish. let _ = py_asyncio_loop_handle.await; @@ -173,9 +102,5 @@ pub(super) async fn start( sig_handle.abort(); } - if let Some(error) = error_slot.lock().unwrap().take() { - return Err(error); - } - result } diff --git a/packages/pynumaflow-lite/src/map/mod.rs b/packages/pynumaflow-lite/src/map/mod.rs index 0232cf00..85c3843c 100644 --- a/packages/pynumaflow-lite/src/map/mod.rs +++ b/packages/pynumaflow-lite/src/map/mod.rs @@ -1,9 +1,6 @@ use std::collections::HashMap; -use std::path::PathBuf; -use std::time::Duration; use numaflow::map; -use numaflow::proto::map::map_client::MapClient; use chrono::{DateTime, Utc}; @@ -11,89 +8,233 @@ use chrono::{DateTime, Utc}; /// and can pass in the Python function. pub mod server; -use tokio::net::UnixStream; -use tokio::time::Instant; -use tonic::transport::Uri; -use tower::service_fn; - use pyo3::prelude::*; use std::sync::Mutex; -pub(crate) fn bytes_literal(value: &[u8]) -> String { - format!("b\"{}\"", String::from_utf8_lossy(value).escape_debug()) +/// SystemMetadata wraps system-generated metadata groups per message. +/// It is read-only to UDFs. +#[pyclass(module = "pynumaflow_lite.mapper", from_py_object)] +#[derive(Clone, Default, Debug)] +pub struct SystemMetadata { + data: HashMap>>, } -fn metadata_literal(metadata: &HashMap>>) -> String { - let groups: Vec = metadata - .iter() - .map(|(group, kv)| { - let entries: Vec = kv - .iter() - .map(|(key, value)| format!("{:?}: {}", key, bytes_literal(value))) - .collect(); - format!("{:?}: {{{}}}", group, entries.join(", ")) - }) - .collect(); - format!("{{{}}}", groups.join(", ")) +#[pymethods] +impl SystemMetadata { + #[new] + #[pyo3(signature = () -> "SystemMetadata")] + fn new() -> Self { + Self::default() + } + + /// Returns the groups of the system metadata. + /// If there are no groups, it returns an empty list. + #[pyo3(signature = () -> "list[str]")] + fn groups(&self) -> Vec { + self.data.keys().cloned().collect() + } + + /// Returns the keys of the system metadata for the given group. + /// If there are no keys or the group is not present, it returns an empty list. + #[pyo3(signature = (group: "str") -> "list[str]")] + fn keys(&self, group: &str) -> Vec { + self.data + .get(group) + .map(|kv| kv.keys().cloned().collect()) + .unwrap_or_default() + } + + /// Returns the value of the system metadata for the given group and key. + /// If there is no value or the group or key is not present, it returns an empty bytes. + #[pyo3(signature = (group: "str", key: "str") -> "bytes")] + fn value(&self, group: &str, key: &str) -> Vec { + self.data + .get(group) + .and_then(|kv| kv.get(key)) + .cloned() + .unwrap_or_default() + } + + fn __repr__(&self) -> String { + format!("SystemMetadata(groups={:?})", self.groups()) + } +} + +impl From for SystemMetadata { + fn from(value: map::SystemMetadata) -> Self { + let mut data = HashMap::new(); + for group in value.groups() { + let mut kv = HashMap::new(); + for key in value.keys(&group) { + kv.insert(key.clone(), value.value(&group, &key)); + } + data.insert(group, kv); + } + Self { data } + } +} + +/// UserMetadata wraps user-defined metadata groups per message. +/// Users can read and write to this metadata. +#[pyclass(module = "pynumaflow_lite.mapper", from_py_object)] +#[derive(Clone, Default, Debug)] +pub struct UserMetadata { + data: HashMap>>, +} + +#[pymethods] +impl UserMetadata { + #[new] + #[pyo3(signature = () -> "UserMetadata")] + fn new() -> Self { + Self::default() + } + + /// Returns the groups of the user metadata. + /// If there are no groups, it returns an empty list. + #[pyo3(signature = () -> "list[str]")] + fn groups(&self) -> Vec { + self.data.keys().cloned().collect() + } + + /// Returns the keys of the user metadata for the given group. + /// If there are no keys or the group is not present, it returns an empty list. + #[pyo3(signature = (group: "str") -> "list[str]")] + fn keys(&self, group: &str) -> Vec { + self.data + .get(group) + .map(|kv| kv.keys().cloned().collect()) + .unwrap_or_default() + } + + /// Returns the value of the user metadata for the given group and key. + /// If there is no value or the group or key is not present, it returns an empty bytes. + #[pyo3(signature = (group: "str", key: "str") -> "bytes")] + fn value(&self, group: &str, key: &str) -> Vec { + self.data + .get(group) + .and_then(|kv| kv.get(key)) + .cloned() + .unwrap_or_default() + } + + /// Creates a new group in the user metadata. + /// If the group already exists, this is a no-op. + #[pyo3(signature = (group: "str"))] + fn create_group(&mut self, group: String) { + self.data.entry(group).or_default(); + } + + /// Adds a key-value pair to the user metadata. + /// If the group is not present, it creates a new group. + #[pyo3(signature = (group: "str", key: "str", value: "bytes"))] + fn add_kv(&mut self, group: String, key: String, value: Vec) { + self.data.entry(group).or_default().insert(key, value); + } + + /// Removes a key from a group in the user metadata. + /// If the key or group is not present, it's a no-op. + #[pyo3(signature = (group: "str", key: "str"))] + fn remove_key(&mut self, group: &str, key: &str) { + if let Some(kv) = self.data.get_mut(group) { + kv.remove(key); + } + } + + /// Removes a group from the user metadata. + /// If the group is not present, it's a no-op. + #[pyo3(signature = (group: "str"))] + fn remove_group(&mut self, group: &str) { + self.data.remove(group); + } + + fn __repr__(&self) -> String { + format!("UserMetadata(groups={:?})", self.groups()) + } } -fn system_metadata_to_hash_map( - value: map::SystemMetadata, -) -> HashMap>> { - let mut data = HashMap::new(); - for group in value.groups() { - let mut kv = HashMap::new(); - for key in value.keys(&group) { - kv.insert(key.clone(), value.value(&group, &key)); +impl From for UserMetadata { + fn from(value: map::UserMetadata) -> Self { + let mut data = HashMap::new(); + for group in value.groups() { + let mut kv = HashMap::new(); + for key in value.keys(&group) { + kv.insert(key.clone(), value.value(&group, &key)); + } + data.insert(group, kv); } - data.insert(group, kv); + Self { data } } - data } -fn user_metadata_to_hash_map( - value: map::UserMetadata, -) -> HashMap>> { - let mut data = HashMap::new(); - for group in value.groups() { - let mut kv = HashMap::new(); - for key in value.keys(&group) { - kv.insert(key.clone(), value.value(&group, &key)); +impl From for map::UserMetadata { + fn from(value: UserMetadata) -> Self { + let mut user_metadata = map::UserMetadata::new(); + for (group, kv_map) in value.data { + for (key, val) in kv_map { + user_metadata.add_kv(group.clone(), key, val); + } } - data.insert(group, kv); + user_metadata + } +} + +/// A collection of [Message]s. +#[pyclass(module = "pynumaflow_lite.mapper", from_py_object)] +#[derive(Clone, Debug)] +pub struct Messages { + pub(crate) messages: Vec, +} + +#[pymethods] +impl Messages { + #[new] + #[pyo3(signature = () -> "Messages")] + fn new() -> Self { + Self { messages: vec![] } + } + + /// Append a [Message] to the collection. + #[pyo3(signature = (message: "Message"))] + fn append(&mut self, message: Message) { + self.messages.push(message); + } + + fn __repr__(&self) -> String { + format!("Messages({:?})", self.messages) + } + + fn __str__(&self) -> String { + format!("Messages({:?})", self.messages) } - data } /// A message to be sent to the next vertex. -#[pyclass(module = "pynumaflow_lite.mapper", from_py_object, eq)] -#[derive(Clone, Default, Debug, PartialEq)] +#[pyclass(module = "pynumaflow_lite.mapper", from_py_object)] +#[derive(Clone, Default, Debug)] pub struct Message { /// Keys are a collection of strings which will be passed on to the next vertex as is. It can /// be an empty collection. - #[pyo3(get)] pub keys: Option>, /// Value is the value passed to the next vertex. - #[pyo3(get)] pub value: Vec, /// Tags are used for [conditional forwarding](https://numaflow.numaproj.io/user-guide/reference/conditional-forwarding/). - #[pyo3(get)] pub tags: Option>, /// User metadata for the message. - #[pyo3(get)] - pub user_metadata: Option>>>, + pub user_metadata: Option, } #[pymethods] impl Message { - /// Create a new Message with the given value. Keys, tags, and user_metadata are optional. + /// Create a new [Message] with the given value, keys, tags, and user_metadata. #[new] - #[pyo3(signature = (value: "bytes", keys: "list[str] | None"=None, tags: "list[str] | None"=None, user_metadata: "dict[str, dict[str, bytes]] | None"=None) -> "Message")] + #[pyo3(signature = (value: "bytes", keys: "list[str] | None"=None, tags: "list[str] | None"=None, user_metadata: "UserMetadata | None"=None) -> "Message" + )] fn new( value: Vec, keys: Option>, tags: Option>, - user_metadata: Option>>>, + user_metadata: Option, ) -> Self { Self { keys, @@ -103,10 +244,10 @@ impl Message { } } - /// A Message marked to be dropped, i.e. not forwarded to the next vertex. + /// Drop a [Message], do not forward to the next vertex. + #[pyo3(signature = ())] #[staticmethod] - #[pyo3(signature = () -> "Message")] - fn to_drop() -> Self { + fn message_to_drop() -> Self { Self { keys: None, value: vec![], @@ -114,22 +255,6 @@ impl Message { user_metadata: None, } } - - fn __repr__(&self) -> String { - format!( - "Message(value={}, keys={}, tags={}, user_metadata={})", - bytes_literal(&self.value), - self.keys - .as_ref() - .map_or_else(|| "None".to_string(), |keys| format!("{keys:?}")), - self.tags - .as_ref() - .map_or_else(|| "None".to_string(), |tags| format!("{tags:?}")), - self.user_metadata - .as_ref() - .map_or_else(|| "None".to_string(), metadata_literal), - ) - } } impl From for map::Message { @@ -138,21 +263,12 @@ impl From for map::Message { keys: value.keys, value: value.value, tags: value.tags, - user_metadata: value.user_metadata.map(|m| { - let mut umd = map::UserMetadata::new(); - for (group, kv) in m { - for (key, val) in kv { - umd.add_kv(group.clone(), key, val); - } - } - umd - }), + user_metadata: value.user_metadata.map(|m| m.into()), } } } -/// The incoming Datum passed to the map handler. It carries the event's keys, value, -/// event_time, watermark, headers, and the user/system metadata. +/// The incoming [MapRequest] accessible in Python function. #[pyclass(module = "pynumaflow_lite.mapper")] pub struct Datum { /// Set of keys in the (key, value) terminology of map/reduce paradigm. @@ -167,162 +283,106 @@ pub struct Datum { pub watermark: DateTime, /// Time of the element as seen at source or aligned after a reduce operation. #[pyo3(get)] - pub event_time: DateTime, + pub eventtime: DateTime, /// Headers for the message. #[pyo3(get)] pub headers: HashMap, /// User metadata for the message. #[pyo3(get)] - pub user_metadata: HashMap>>, + pub user_metadata: UserMetadata, /// System metadata for the message. #[pyo3(get)] - pub system_metadata: HashMap>>, + pub system_metadata: SystemMetadata, } -#[pymethods] impl Datum { - #[new] - #[pyo3(signature = ( - *, - keys: "list[str] | None"=None, - value: "bytes | None"=None, - event_time: "datetime.datetime | None"=None, - watermark: "datetime.datetime | None"=None, - headers: "dict[str, str] | None"=None, - user_metadata: "dict[str, dict[str, bytes]] | None"=None, - system_metadata: "dict[str, dict[str, bytes]] | None"=None, - ) -> "Datum")] fn new( - keys: Option>, - value: Option>, - event_time: Option>, - watermark: Option>, - headers: Option>, - user_metadata: Option>>>, - system_metadata: Option>>>, + keys: Vec, + value: Vec, + watermark: DateTime, + eventtime: DateTime, + headers: HashMap, + user_metadata: UserMetadata, + system_metadata: SystemMetadata, ) -> Self { Self { - keys: keys.unwrap_or_default(), - value: value.unwrap_or_default(), - watermark: watermark.unwrap_or(DateTime::::UNIX_EPOCH), - event_time: event_time.unwrap_or(DateTime::::UNIX_EPOCH), - headers: headers.unwrap_or_default(), - user_metadata: user_metadata.unwrap_or_default(), - system_metadata: system_metadata.unwrap_or_default(), + keys, + value, + watermark, + eventtime, + headers, + user_metadata, + system_metadata, } } fn __repr__(&self) -> String { format!( - "Datum(keys={:?}, value={}, watermark={}, event_time={}, headers={:?}, user_metadata={}, system_metadata={})", + "Datum(keys={:?}, value={:?}, watermark={}, eventtime={}, headers={:?}, user_metadata={:?}, system_metadata={:?})", self.keys, - bytes_literal(&self.value), + self.value, self.watermark, - self.event_time, + self.eventtime, self.headers, - metadata_literal(&self.user_metadata), - metadata_literal(&self.system_metadata) + self.user_metadata, + self.system_metadata ) } fn __str__(&self) -> String { - self.__repr__() + format!( + "Datum(keys={:?}, value={:?}, watermark={}, eventtime={}, headers={:?}, user_metadata={:?}, system_metadata={:?})", + self.keys, + String::from_utf8_lossy(&self.value), + self.watermark, + self.eventtime, + self.headers, + self.user_metadata, + self.system_metadata + ) } } impl From for Datum { fn from(value: map::MapRequest) -> Self { - Self { - keys: value.keys, - value: value.value, - watermark: value.watermark, - event_time: value.eventtime, - headers: value.headers, - user_metadata: user_metadata_to_hash_map(value.user_metadata), - system_metadata: system_metadata_to_hash_map(value.system_metadata), - } - } -} - -pub(crate) async fn map_grpc_client( - sock_file: String, -) -> PyResult> { - let endpoint = tonic::transport::Endpoint::try_from("http://[::]:50051") - .map_err(|e| pyo3::PyErr::new::(e.to_string()))?; - - let channel = endpoint - .connect_with_connector(service_fn(move |_: Uri| { - let sock = PathBuf::from(sock_file.clone()); - async move { - Ok::<_, std::io::Error>(hyper_util::rt::TokioIo::new( - UnixStream::connect(sock).await?, - )) - } - })) - .await - .map_err(|e| pyo3::PyErr::new::(e.to_string()))?; - - Ok(MapClient::new(channel)) -} - -pub(crate) async fn wait_for_ready( - sock_file: String, - timeout: Duration, - component: &'static str, -) -> PyResult<()> { - let deadline = Instant::now() + timeout; - - loop { - if let Ok(mut client) = map_grpc_client(sock_file.clone()).await - && let Ok(response) = client.is_ready(()).await - && response.into_inner().ready - { - return Ok(()); - } - - let now = Instant::now(); - if now >= deadline { - return Err(pyo3::PyErr::new::( - format!("timed out waiting for {component} server readiness"), - )); - } - - tokio::time::sleep(std::cmp::min( - Duration::from_millis(100), - deadline.saturating_duration_since(now), - )) - .await; + Datum::new( + value.keys, + value.value, + value.watermark, + value.eventtime, + value.headers, + value.user_metadata.into(), + value.system_metadata.into(), + ) } } /// Async Map Server that can be started from Python code which will run the Python UDF function. -#[pyclass(name = "_MapAsyncServer", module = "pynumaflow_lite.mapper")] +#[pyclass(module = "pynumaflow_lite.mapper")] pub struct MapAsyncServer { sock_file: String, - server_info_file: String, + info_file: String, shutdown_tx: Mutex>>, } #[pymethods] impl MapAsyncServer { #[new] - #[pyo3(signature = ( - sock_file: "str | None"=None, - server_info_file: "str | None"=None, - ) -> "_MapAsyncServer")] - fn new(sock_file: Option, server_info_file: Option) -> Self { + #[pyo3(signature = (sock_file: "str | None"=map::SOCK_ADDR.to_string(), info_file: "str | None"=map::SERVER_INFO_FILE.to_string()) -> "MapAsyncServer" + )] + fn new(sock_file: String, info_file: String) -> Self { Self { - sock_file: sock_file.unwrap_or_else(|| map::SOCK_ADDR.to_string()), - server_info_file: server_info_file.unwrap_or_else(|| map::SERVER_INFO_FILE.to_string()), + sock_file, + info_file, shutdown_tx: Mutex::new(None), } } /// Start the server with the given Python function. - #[pyo3(signature = (handler: "callable") -> "None")] - pub fn start<'a>(&self, py: Python<'a>, handler: Py) -> PyResult> { + #[pyo3(signature = (py_func: "callable") -> "None")] + pub fn start<'a>(&self, py: Python<'a>, py_func: Py) -> PyResult> { let sock_file = self.sock_file.clone(); - let server_info_file = self.server_info_file.clone(); + let info_file = self.info_file.clone(); let (tx, rx) = tokio::sync::oneshot::channel::<()>(); { let mut guard = self.shutdown_tx.lock().unwrap(); @@ -330,25 +390,9 @@ impl MapAsyncServer { } pyo3_async_runtimes::tokio::future_into_py(py, async move { - crate::map::server::start(handler, sock_file, server_info_file, rx).await?; - Ok(()) - }) - } - - /// Wait until the Numaflow IsReady probe succeeds over the map UDS. - #[pyo3(signature = (timeout: "float"=30.0) -> "None")] - pub fn wait_ready<'a>(&self, py: Python<'a>, timeout: f64) -> PyResult> { - if !timeout.is_finite() || timeout < 0.0 { - return Err(pyo3::PyErr::new::( - "timeout must be a non-negative finite float", - )); - } - - let sock_file = self.sock_file.clone(); - let timeout = Duration::from_secs_f64(timeout); - - pyo3_async_runtimes::tokio::future_into_py(py, async move { - wait_for_ready(sock_file, timeout, "map").await?; + crate::map::server::start(py_func, sock_file, info_file, rx) + .await + .expect("server failed to start"); Ok(()) }) } @@ -365,6 +409,9 @@ impl MapAsyncServer { /// Helper to populate a PyModule with map types/functions. pub(crate) fn populate_py_module(m: &Bound) -> PyResult<()> { + m.add_class::()?; + m.add_class::()?; + m.add_class::()?; m.add_class::()?; m.add_class::()?; m.add_class::()?; diff --git a/packages/pynumaflow-lite/src/map/server.rs b/packages/pynumaflow-lite/src/map/server.rs index 98a7a194..1e220994 100644 --- a/packages/pynumaflow-lite/src/map/server.rs +++ b/packages/pynumaflow-lite/src/map/server.rs @@ -1,83 +1,40 @@ +use crate::map::{Datum, Messages}; use numaflow::map; use numaflow::shared::ServerExtras; -use pyo3::exceptions::PyTypeError; use pyo3::prelude::*; -use std::sync::{Arc, Mutex}; +use std::sync::Arc; pub(crate) struct PyMapRunner { pub(crate) event_loop: Arc>, pub(crate) py_func: Arc>, - pub(crate) error_slot: Arc>>, - pub(crate) shutdown_tx: Arc>>>, -} - -impl PyMapRunner { - fn fail(&self, error: PyErr) -> Vec { - // Only the first error is reported; later requests may still be in flight - // while shutdown is underway, and their failures would be duplicates. - let mut error_slot = self.error_slot.lock().unwrap(); - if error_slot.is_none() { - Python::attach(|py| error.print(py)); - *error_slot = Some(error); - } - drop(error_slot); - - if let Some(tx) = self.shutdown_tx.lock().unwrap().take() { - let _ = tx.send(()); - } - - Vec::new() - } } #[tonic::async_trait] impl map::Mapper for PyMapRunner { async fn map(&self, input: map::MapRequest) -> Vec { - let fut = match Python::attach(|py| -> PyResult<_> { + let fut = Python::attach(|py| { + let keys = input.keys.clone(); + let input: Datum = input.into(); + let py_func = self.py_func.clone(); + let locals = pyo3_async_runtimes::TaskLocals::new(self.event_loop.bind(py).clone()); - let datum: crate::map::Datum = input.into(); - let coro = self.py_func.call1(py, (datum,))?.into_bound(py); - let is_awaitable: bool = py - .import("inspect")? - .call_method1("isawaitable", (&coro,))? - .extract()?; - if !is_awaitable { - return Err(PyErr::new::( - "map handler must be an async function (coroutine)", - )); - } - pyo3_async_runtimes::into_future_with_locals(&locals, coro).map_err(|_| { - PyErr::new::("map handler must be an async function (coroutine)") - }) - }) { - Ok(fut) => fut, - Err(error) => return self.fail(error), - }; - - let result = match fut.await { - Ok(result) => result, - Err(error) => return self.fail(error), - }; - - let messages: Vec = match Python::attach(|py| { - result.extract(py).map_err(|_| { - let type_name = result - .bind(py) - .get_type() - .name() - .map(|name| name.to_string_lossy().into_owned()) - .unwrap_or_else(|_| "".to_string()); - PyErr::new::(format!( - "map handler must return list[Message], got {type_name}" - )) - }) - }) { - Ok(messages) => messages, - Err(error) => return self.fail(error), - }; - - messages.into_iter().map(|m| m.into()).collect() + + let coro = py_func.call1(py, (keys, input)).unwrap().into_bound(py); + + pyo3_async_runtimes::into_future_with_locals(&locals, coro).unwrap() + }); + + let result = fut.await.unwrap(); + + let result = Python::attach(|py| { + let x: Messages = result.extract(py).unwrap(); + x + }); + + println!("{:?}", result); + + result.messages.into_iter().map(|m| m.into()).collect() } } @@ -89,34 +46,14 @@ pub(super) async fn start( shutdown_rx: tokio::sync::oneshot::Receiver<()>, ) -> Result<(), pyo3::PyErr> { let (tx, rx) = tokio::sync::oneshot::channel(); - let py_asyncio_loop_handle = tokio::task::spawn_blocking({ - println!( - "Starting Map UDF. socket={}, server_info={}", - sock_file, info_file - ); - move || crate::pyrs::run_asyncio(tx) - }); + let py_asyncio_loop_handle = tokio::task::spawn_blocking(move || crate::pyrs::run_asyncio(tx)); let event_loop = rx.await.unwrap(); - let error_slot = Arc::new(Mutex::new(None)); - let (internal_shutdown_tx, internal_shutdown_rx) = tokio::sync::oneshot::channel(); - let (server_shutdown_tx, server_shutdown_rx) = tokio::sync::oneshot::channel(); - - tokio::spawn(async move { - tokio::select! { - _ = shutdown_rx => {}, - _ = internal_shutdown_rx => {}, - } - let _ = server_shutdown_tx.send(()); - }); - - let (sig_handle, combined_rx) = crate::pyrs::setup_sig_handler(server_shutdown_rx); + let (sig_handle, combined_rx) = crate::pyrs::setup_sig_handler(shutdown_rx); let py_map_runner = PyMapRunner { py_func: Arc::new(py_func), event_loop: event_loop.clone(), - error_slot: error_slot.clone(), - shutdown_tx: Arc::new(Mutex::new(Some(internal_shutdown_tx))), }; let server = numaflow::map::Server::new(py_map_runner) @@ -135,7 +72,7 @@ pub(super) async fn start( } }); - println!("Numaflow Map has shutdown..."); + println!("Numaflow Core has shutdown..."); // Wait for the blocking asyncio thread to finish. let _ = py_asyncio_loop_handle.await; @@ -146,9 +83,5 @@ pub(super) async fn start( sig_handle.abort(); } - if let Some(error) = error_slot.lock().unwrap().take() { - return Err(error); - } - result } diff --git a/packages/pynumaflow-lite/src/mapstream/mod.rs b/packages/pynumaflow-lite/src/mapstream/mod.rs index 84288c57..24e126f0 100644 --- a/packages/pynumaflow-lite/src/mapstream/mod.rs +++ b/packages/pynumaflow-lite/src/mapstream/mod.rs @@ -2,7 +2,6 @@ use chrono::{DateTime, Utc}; use numaflow::mapstream; use std::collections::HashMap; use std::sync::Mutex; -use std::time::Duration; pub mod server; @@ -10,7 +9,8 @@ pub mod server; use pyo3::prelude::*; /// Streaming Datum mirrors MapStreamRequest for Python -#[pyclass(module = "pynumaflow_lite.mapstreamer")] +#[pyclass(module = "pynumaflow_lite.mapstreamer", from_py_object)] +#[derive(Clone)] pub struct Datum { /// Set of keys in the (key, value) terminology of the map/reduce paradigm. #[pyo3(get)] @@ -24,76 +24,33 @@ pub struct Datum { pub watermark: DateTime, /// Time of the element as seen at source or aligned after a reduce operation. #[pyo3(get)] - pub event_time: DateTime, + pub eventtime: DateTime, /// Headers associated with the message. #[pyo3(get)] pub headers: HashMap, } impl Datum { - pub(crate) fn from_parts( + pub(crate) fn new( keys: Vec, value: Vec, watermark: DateTime, - event_time: DateTime, + eventtime: DateTime, headers: HashMap, ) -> Self { Self { keys, value, watermark, - event_time, + eventtime, headers, } } } -#[pymethods] -impl Datum { - #[new] - #[pyo3(signature = ( - *, - keys: "list[str] | None"=None, - value: "bytes | None"=None, - event_time: "datetime.datetime | None"=None, - watermark: "datetime.datetime | None"=None, - headers: "dict[str, str] | None"=None, - ) -> "Datum")] - fn new( - keys: Option>, - value: Option>, - event_time: Option>, - watermark: Option>, - headers: Option>, - ) -> Self { - Self { - keys: keys.unwrap_or_default(), - value: value.unwrap_or_default(), - watermark: watermark.unwrap_or(DateTime::::UNIX_EPOCH), - event_time: event_time.unwrap_or(DateTime::::UNIX_EPOCH), - headers: headers.unwrap_or_default(), - } - } - - fn __repr__(&self) -> String { - format!( - "Datum(keys={:?}, value={}, watermark={}, event_time={}, headers={:?})", - self.keys, - crate::map::bytes_literal(&self.value), - self.watermark, - self.event_time, - self.headers - ) - } - - fn __str__(&self) -> String { - self.__repr__() - } -} - impl From for Datum { fn from(value: numaflow::mapstream::MapStreamRequest) -> Self { - Self::from_parts( + Self::new( value.keys, value.value, value.watermark, @@ -104,17 +61,14 @@ impl From for Datum { } /// A message to be sent downstream from a streaming handler. -#[pyclass(module = "pynumaflow_lite.mapstreamer", from_py_object, eq)] -#[derive(Clone, Default, Debug, PartialEq)] +#[pyclass(module = "pynumaflow_lite.mapstreamer", from_py_object)] +#[derive(Clone, Default, Debug)] pub struct Message { /// Keys are a collection of strings which will be passed on to the next vertex as is. - #[pyo3(get)] pub keys: Option>, /// Value is the value passed to the next vertex. - #[pyo3(get)] pub value: Vec, /// Tags are used for conditional forwarding. - #[pyo3(get)] pub tags: Option>, } @@ -127,10 +81,10 @@ impl Message { Self { keys, value, tags } } - /// A Message marked to be dropped, i.e. not forwarded to the next vertex. + /// Drop a Message, do not forward to the next vertex. + #[pyo3(signature = ())] #[staticmethod] - #[pyo3(signature = () -> "Message")] - fn to_drop() -> Self { + fn message_to_drop() -> Self { Self { keys: None, value: vec![], @@ -138,17 +92,11 @@ impl Message { } } - fn __repr__(&self) -> String { - format!( - "Message(value={}, keys={}, tags={})", - crate::map::bytes_literal(&self.value), - self.keys - .as_ref() - .map_or_else(|| "None".to_string(), |keys| format!("{keys:?}")), - self.tags - .as_ref() - .map_or_else(|| "None".to_string(), |tags| format!("{tags:?}")), - ) + /// Convenience alias to match example usage: Message.to_drop() + #[pyo3(signature = ())] + #[staticmethod] + fn to_drop() -> Self { + Self::message_to_drop() } } @@ -163,34 +111,31 @@ impl From for mapstream::Message { } /// Async MapStream Server that can be started from Python code which will run the Python UDF async generator. -#[pyclass(name = "_MapStreamAsyncServer", module = "pynumaflow_lite.mapstreamer")] +#[pyclass(module = "pynumaflow_lite.mapstreamer")] pub struct MapStreamAsyncServer { sock_file: String, - server_info_file: String, + info_file: String, shutdown_tx: Mutex>>, } #[pymethods] impl MapStreamAsyncServer { #[new] - #[pyo3(signature = ( - sock_file: "str | None"=None, - server_info_file: "str | None"=None, - ) -> "_MapStreamAsyncServer")] - fn new(sock_file: Option, server_info_file: Option) -> Self { + #[pyo3(signature = (sock_file: "str | None"=mapstream::SOCK_ADDR.to_string(), info_file: "str | None"=mapstream::SERVER_INFO_FILE.to_string()) -> "MapStreamAsyncServer" + )] + fn new(sock_file: String, info_file: String) -> Self { Self { - sock_file: sock_file.unwrap_or_else(|| mapstream::SOCK_ADDR.to_string()), - server_info_file: server_info_file - .unwrap_or_else(|| mapstream::SERVER_INFO_FILE.to_string()), + sock_file, + info_file, shutdown_tx: Mutex::new(None), } } /// Start the server with the given Python async generator function. - #[pyo3(signature = (handler: "callable") -> "None")] - pub fn start<'a>(&self, py: Python<'a>, handler: Py) -> PyResult> { + #[pyo3(signature = (py_func: "callable") -> "None")] + pub fn start<'a>(&self, py: Python<'a>, py_func: Py) -> PyResult> { let sock_file = self.sock_file.clone(); - let server_info_file = self.server_info_file.clone(); + let info_file = self.info_file.clone(); let (tx, rx) = tokio::sync::oneshot::channel::<()>(); { let mut guard = self.shutdown_tx.lock().unwrap(); @@ -198,25 +143,9 @@ impl MapStreamAsyncServer { } pyo3_async_runtimes::tokio::future_into_py(py, async move { - crate::mapstream::server::start(handler, sock_file, server_info_file, rx).await?; - Ok(()) - }) - } - - /// Wait until the Numaflow IsReady probe succeeds over the mapstream UDS. - #[pyo3(signature = (timeout: "float"=30.0) -> "None")] - pub fn wait_ready<'a>(&self, py: Python<'a>, timeout: f64) -> PyResult> { - if !timeout.is_finite() || timeout < 0.0 { - return Err(pyo3::PyErr::new::( - "timeout must be a non-negative finite float", - )); - } - - let sock_file = self.sock_file.clone(); - let timeout = Duration::from_secs_f64(timeout); - - pyo3_async_runtimes::tokio::future_into_py(py, async move { - crate::map::wait_for_ready(sock_file, timeout, "mapstream").await?; + crate::mapstream::server::start(py_func, sock_file, info_file, rx) + .await + .expect("server failed to start"); Ok(()) }) } diff --git a/packages/pynumaflow-lite/src/mapstream/server.rs b/packages/pynumaflow-lite/src/mapstream/server.rs index 2edde058..7dbcaee1 100644 --- a/packages/pynumaflow-lite/src/mapstream/server.rs +++ b/packages/pynumaflow-lite/src/mapstream/server.rs @@ -5,65 +5,34 @@ use crate::pyiterables::PyAsyncIterStream; use numaflow::mapstream; use numaflow::shared::ServerExtras; -use pyo3::exceptions::PyTypeError; use pyo3::prelude::*; -use std::sync::{Arc, Mutex}; +use std::sync::Arc; use tokio::sync::mpsc::Sender; use tokio_stream::StreamExt; pub(crate) struct PyMapStreamRunner { pub(crate) event_loop: Arc>, pub(crate) py_func: Arc>, - pub(crate) error_slot: Arc>>, - pub(crate) shutdown_tx: Arc>>>, -} - -impl PyMapStreamRunner { - fn fail(&self, error: PyErr) { - let mut error_slot = self.error_slot.lock().unwrap(); - if error_slot.is_none() { - Python::attach(|py| error.print(py)); - *error_slot = Some(error); - } - drop(error_slot); - - if let Some(tx) = self.shutdown_tx.lock().unwrap().take() { - let _ = tx.send(()); - } - } } #[tonic::async_trait] impl mapstream::MapStreamer for PyMapStreamRunner { async fn map_stream(&self, input: mapstream::MapStreamRequest, tx: Sender) { - // Call Python handler: handler(datum) -> AsyncIterable[Message] - let agen_obj = match Python::attach(|py| -> PyResult> { + // Call Python handler: handler(keys, datum) -> AsyncIterator + let agen_obj = Python::attach(|py| { + let keys = input.keys.clone(); let datum: Datum = input.into(); let py_func = self.py_func.clone(); - let agen = py_func.call1(py, (datum,))?; - if !agen.bind(py).hasattr("__aiter__")? { - return Err(PyErr::new::( - "mapstream handler must return an async iterable of Message", - )); - } - Ok(agen) - }) { - Ok(agen_obj) => agen_obj, - Err(error) => { - self.fail(error); - return; - } - }; + let agen = py_func + .call1(py, (keys, datum)) + .expect("python handler raised before returning async iterable"); + // Keep as Py + agen.clone_ref(py).extract(py).unwrap_or(agen) + }); // Wrap the Python AsyncIterable in a Rust Stream that yields incrementally - let mut stream = - match PyAsyncIterStream::::new(agen_obj, self.event_loop.clone()) { - Ok(stream) => stream, - Err(error) => { - self.fail(error); - return; - } - }; + let mut stream = PyAsyncIterStream::::new(agen_obj, self.event_loop.clone()) + .expect("failed to construct PyAsyncIterStream"); // Forward each yielded message immediately to the sender while let Some(item) = stream.next().await { @@ -75,7 +44,8 @@ impl mapstream::MapStreamer for PyMapStreamRunner { } } Err(e) => { - self.fail(e); + // Non-stop errors are surfaced per-item; log and stop this stream. + eprintln!("Python async iteration error: {:?}", e); break; } } @@ -87,43 +57,23 @@ impl mapstream::MapStreamer for PyMapStreamRunner { pub(super) async fn start( py_func: Py, sock_file: String, - server_info_file: String, + info_file: String, shutdown_rx: tokio::sync::oneshot::Receiver<()>, ) -> Result<(), pyo3::PyErr> { let (tx, rx) = tokio::sync::oneshot::channel(); - let py_asyncio_loop_handle = tokio::task::spawn_blocking({ - println!( - "Starting MapStream UDF. socket={}, server_info={}", - sock_file, server_info_file - ); - move || crate::pyrs::run_asyncio(tx) - }); + let py_asyncio_loop_handle = tokio::task::spawn_blocking(move || crate::pyrs::run_asyncio(tx)); let event_loop = rx.await.unwrap(); - let error_slot = Arc::new(Mutex::new(None)); - let (internal_shutdown_tx, internal_shutdown_rx) = tokio::sync::oneshot::channel(); - let (server_shutdown_tx, server_shutdown_rx) = tokio::sync::oneshot::channel(); - - tokio::spawn(async move { - tokio::select! { - _ = shutdown_rx => {}, - _ = internal_shutdown_rx => {}, - } - let _ = server_shutdown_tx.send(()); - }); - - let (sig_handle, combined_rx) = crate::pyrs::setup_sig_handler(server_shutdown_rx); + let (sig_handle, combined_rx) = crate::pyrs::setup_sig_handler(shutdown_rx); let py_runner = PyMapStreamRunner { py_func: Arc::new(py_func), event_loop: event_loop.clone(), - error_slot: error_slot.clone(), - shutdown_tx: Arc::new(Mutex::new(Some(internal_shutdown_tx))), }; let server = numaflow::mapstream::Server::new(py_runner) .with_socket_file(sock_file) - .with_server_info_file(server_info_file); + .with_server_info_file(info_file); let result = server .start_with_shutdown(combined_rx) @@ -137,7 +87,7 @@ pub(super) async fn start( } }); - println!("Numaflow MapStream has shutdown..."); + println!("Numaflow Core (stream) has shutdown..."); // Wait for the blocking asyncio thread to finish. let _ = py_asyncio_loop_handle.await; @@ -148,9 +98,5 @@ pub(super) async fn start( sig_handle.abort(); } - if let Some(error) = error_slot.lock().unwrap().take() { - return Err(error); - } - result } diff --git a/packages/pynumaflow-lite/src/sink/server.rs b/packages/pynumaflow-lite/src/sink/server.rs index f588bdf7..671ac02f 100644 --- a/packages/pynumaflow-lite/src/sink/server.rs +++ b/packages/pynumaflow-lite/src/sink/server.rs @@ -125,7 +125,7 @@ pub(super) async fn start( let py_asyncio_loop_handle = tokio::task::spawn_blocking({ println!( "Starting Sink UDF. socket={}, server_info={}", - sock_file, info_file + &sock_file, &info_file ); move || crate::pyrs::run_asyncio(tx) }); diff --git a/packages/pynumaflow-lite/tests/_test_utils.py b/packages/pynumaflow-lite/tests/_test_utils.py index 3d272280..fced4816 100644 --- a/packages/pynumaflow-lite/tests/_test_utils.py +++ b/packages/pynumaflow-lite/tests/_test_utils.py @@ -37,16 +37,6 @@ async def _unused_handler(datums): asyncio.run(server.wait_ready(timeout=timeout)) -def _wait_for_map_ready(path: Path, timeout: float = 10.0) -> None: - from pynumaflow_lite import mapper - - async def _unused_handler(_datum): - return [] - - server = mapper.MapAsyncServer(_unused_handler, sock_file=str(path)) - asyncio.run(server.wait_ready(timeout=timeout)) - - def run_python_server_with_rust_client( script: str, sock_path: Path, @@ -101,8 +91,6 @@ def run_python_server_with_rust_client( try: if rust_bin_name == "test_sink": _wait_for_sink_ready(sock_path, timeout=socket_timeout) - elif rust_bin_name == "test_map": - _wait_for_map_ready(sock_path, timeout=socket_timeout) else: _wait_for_socket(sock_path, timeout=socket_timeout) diff --git a/packages/pynumaflow-lite/tests/examples/batchmap_cat.py b/packages/pynumaflow-lite/tests/examples/batchmap_cat.py index 1558ca73..64a018e8 100644 --- a/packages/pynumaflow-lite/tests/examples/batchmap_cat.py +++ b/packages/pynumaflow-lite/tests/examples/batchmap_cat.py @@ -1,24 +1,51 @@ -from collections.abc import AsyncIterator +import asyncio +import collections.abc +import signal +from collections.abc import Awaitable, Callable from pynumaflow_lite import batchmapper -from pynumaflow_lite.batchmapper import BatchResponse, Datum, Message -async def async_handler(batch: AsyncIterator[Datum]) -> list[BatchResponse]: - responses = [] +async def async_handler( + batch: collections.abc.AsyncIterator[batchmapper.Datum], +) -> batchmapper.BatchResponses: + responses = batchmapper.BatchResponses() async for d in batch: - resp = BatchResponse(d.id) + resp = batchmapper.BatchResponse.from_id(d.id) if d.value == b"bad world": - resp.append(Message.to_drop()) - else: - resp.append(Message(d.value, keys=d.keys)) + resp.append(batchmapper.Message.message_to_drop()) + continue + + resp.append(batchmapper.Message(d.value, d.keys)) responses.append(resp) return responses +async def start( + f: Callable[ + [collections.abc.AsyncIterator[batchmapper.Datum]], + Awaitable[batchmapper.BatchResponses], + ], +): + sock_file = "/tmp/var/run/numaflow/batchmap.sock" + server_info_file = "/tmp/var/run/numaflow/mapper-server-info" + server = batchmapper.BatchMapAsyncServer(sock_file, server_info_file) + + # Register loop-level signal handlers to request graceful shutdown + 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 + + if __name__ == "__main__": - batchmapper.BatchMapAsyncServer( - async_handler, - sock_file="/tmp/var/run/numaflow/batchmap.sock", - server_info_file="/tmp/var/run/numaflow/mapper-server-info", - ).run() + asyncio.run(start(async_handler)) diff --git a/packages/pynumaflow-lite/tests/examples/batchmap_cat_class.py b/packages/pynumaflow-lite/tests/examples/batchmap_cat_class.py index e3dd6dd1..f78cfdfe 100644 --- a/packages/pynumaflow-lite/tests/examples/batchmap_cat_class.py +++ b/packages/pynumaflow-lite/tests/examples/batchmap_cat_class.py @@ -1,24 +1,48 @@ -from collections.abc import AsyncIterator +import asyncio +import signal +from collections.abc import AsyncIterator, Awaitable, Callable from pynumaflow_lite import batchmapper -from pynumaflow_lite.batchmapper import BatchResponse, Datum, Message +from pynumaflow_lite.batchmapper import Message -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 - ] +class SimpleBatchCat(batchmapper.BatchMapper): + async def handler(self, batch: AsyncIterator[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[[AsyncIterator[batchmapper.Datum]], Awaitable[batchmapper.BatchResponses]], +): + sock_file = "/tmp/var/run/numaflow/batchmap.sock" + server_info_file = "/tmp/var/run/numaflow/mapper-server-info" + server = batchmapper.BatchMapAsyncServer(sock_file, server_info_file) + + # 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 if __name__ == "__main__": - batch_mapper_obj = SimpleBatchCat() - batchmapper.BatchMapAsyncServer( - batch_mapper_obj.handler, - sock_file="/tmp/var/run/numaflow/batchmap.sock", - server_info_file="/tmp/var/run/numaflow/mapper-server-info", - ).run() + async_handler = SimpleBatchCat() + asyncio.run(start(async_handler)) diff --git a/packages/pynumaflow-lite/tests/examples/map_cat.py b/packages/pynumaflow-lite/tests/examples/map_cat.py index f7947540..9e04728b 100644 --- a/packages/pynumaflow-lite/tests/examples/map_cat.py +++ b/packages/pynumaflow-lite/tests/examples/map_cat.py @@ -1,34 +1,61 @@ +import asyncio +import signal +from collections.abc import Awaitable, Callable + from pynumaflow_lite import mapper -async def async_handler(datum: mapper.Datum) -> list[mapper.Message]: +async def async_handler(keys: list[str], payload: mapper.Datum) -> mapper.Messages: + messages = mapper.Messages() + # Read system metadata (read-only) - print(f"System metadata groups: {list(datum.system_metadata)}") - for group, key_values in datum.system_metadata.items(): - for key, value in key_values.items(): + print(f"System metadata groups: {payload.system_metadata.groups()}") + for group in payload.system_metadata.groups(): + for key in payload.system_metadata.keys(group): + value = payload.system_metadata.value(group, key) print(f" System[{group}][{key}] = {value}") - # Read user metadata - print(f"User metadata groups: {list(datum.user_metadata)}") - for group, key_values in datum.user_metadata.items(): - for key, value in key_values.items(): + # Read user metadata (read-only from input) + print(f"User metadata groups: {payload.user_metadata.groups()}") + for group in payload.user_metadata.groups(): + for key in payload.user_metadata.keys(group): + value = payload.user_metadata.value(group, key) print(f" User[{group}][{key}] = {value}") - if datum.value == b"bad world": - return [mapper.Message.to_drop()] + if payload.value == b"bad world": + messages.append(mapper.Message.message_to_drop()) + else: + # Create user metadata for the outgoing message + user_metadata = mapper.UserMetadata() + user_metadata.create_group("processing") + user_metadata.add_kv("processing", "handler", b"map_cat") + user_metadata.add_kv("processing", "msg_length", str(len(payload.value)).encode()) + + messages.append(mapper.Message(payload.value, keys, user_metadata=user_metadata)) + + return messages + + +async def start(f: Callable[[list[str], mapper.Datum], Awaitable[mapper.Messages]]): + sock_file = "/tmp/var/run/numaflow/map.sock" + server_info_file = "/tmp/var/run/numaflow/mapper-server-info" + server = mapper.MapAsyncServer(sock_file, server_info_file) + + # Register loop-level signal handlers to request graceful shutdown + 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 - user_metadata = { - "processing": { - "handler": b"map_cat", - "msg_length": str(len(datum.value)).encode(), - } - } - return [mapper.Message(datum.value, keys=datum.keys, user_metadata=user_metadata)] + try: + await server.start(f) + print("Shutting down gracefully...") + except asyncio.CancelledError: + server.stop() + return if __name__ == "__main__": - mapper.MapAsyncServer( - async_handler, - sock_file="/tmp/var/run/numaflow/map.sock", - server_info_file="/tmp/var/run/numaflow/mapper-server-info", - ).run() + asyncio.run(start(async_handler)) diff --git a/packages/pynumaflow-lite/tests/examples/map_cat_class.py b/packages/pynumaflow-lite/tests/examples/map_cat_class.py index b9b76ec6..079beaf2 100644 --- a/packages/pynumaflow-lite/tests/examples/map_cat_class.py +++ b/packages/pynumaflow-lite/tests/examples/map_cat_class.py @@ -1,36 +1,70 @@ +import asyncio +import signal +from collections.abc import Awaitable, Callable + from pynumaflow_lite import mapper -class SimpleCat: - async def handler(self, datum: mapper.Datum) -> list[mapper.Message]: +class SimpleCat(mapper.Mapper): + async def handler(self, keys: list[str], payload: mapper.Datum) -> mapper.Messages: + + messages = mapper.Messages() + # Read system metadata (read-only) - print(f"System metadata groups: {list(datum.system_metadata)}") - for group, key_values in datum.system_metadata.items(): - for key, value in key_values.items(): + print(f"System metadata groups: {payload.system_metadata.groups()}") + for group in payload.system_metadata.groups(): + for key in payload.system_metadata.keys(group): + value = payload.system_metadata.value(group, key) print(f" System[{group}][{key}] = {value}") - # Read user metadata - print(f"User metadata groups: {list(datum.user_metadata)}") - for group, key_values in datum.user_metadata.items(): - for key, value in key_values.items(): + # Read user metadata (read-only from input) + print(f"User metadata groups: {payload.user_metadata.groups()}") + for group in payload.user_metadata.groups(): + for key in payload.user_metadata.keys(group): + value = payload.user_metadata.value(group, key) print(f" User[{group}][{key}] = {value}") - if datum.value == b"bad world": - return [mapper.Message.to_drop()] + if payload.value == b"bad world": + messages.append(mapper.Message.message_to_drop()) + else: + # Create user metadata for the outgoing message + user_metadata = mapper.UserMetadata() + user_metadata.create_group("processing") + user_metadata.add_kv("processing", "handler", b"map_cat_class") + user_metadata.add_kv("processing", "msg_length", str(len(payload.value)).encode()) + + messages.append(mapper.Message(payload.value, keys, user_metadata=user_metadata)) + + return messages + + +async def start(f: Callable[[list[str], mapper.Datum], Awaitable[mapper.Messages]]): + sock_file = "/tmp/var/run/numaflow/map.sock" + server_info_file = "/tmp/var/run/numaflow/mapper-server-info" + server = mapper.MapAsyncServer(sock_file, server_info_file) + + # 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 - user_metadata = { - "processing": { - "handler": b"map_cat_class", - "msg_length": str(len(datum.value)).encode(), - } - } - return [mapper.Message(datum.value, keys=datum.keys, user_metadata=user_metadata)] + 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 if __name__ == "__main__": - mapper_obj = SimpleCat() - mapper.MapAsyncServer( - mapper_obj.handler, - sock_file="/tmp/var/run/numaflow/map.sock", - server_info_file="/tmp/var/run/numaflow/mapper-server-info", - ).run() + async_handler = SimpleCat() + asyncio.run(start(async_handler)) diff --git a/packages/pynumaflow-lite/tests/examples/mapstream_cat.py b/packages/pynumaflow-lite/tests/examples/mapstream_cat.py index d809d7e3..0fd0bcbb 100644 --- a/packages/pynumaflow-lite/tests/examples/mapstream_cat.py +++ b/packages/pynumaflow-lite/tests/examples/mapstream_cat.py @@ -1,10 +1,12 @@ -from collections.abc import AsyncIterable +import asyncio +import signal +from collections.abc import AsyncIterator, Callable from pynumaflow_lite import mapstreamer from pynumaflow_lite.mapstreamer import Message -async def async_handler(datum: mapstreamer.Datum) -> AsyncIterable[Message]: +async def async_handler(keys: list[str], datum: mapstreamer.Datum) -> AsyncIterator[Message]: """ A handler that splits the input datum value into multiple strings by `,` separator and emits them as a stream. @@ -14,12 +16,29 @@ async def async_handler(datum: mapstreamer.Datum) -> AsyncIterable[Message]: yield Message.to_drop() return for s in parts: - yield Message(s.encode(), keys=datum.keys) + yield Message(s.encode(), keys) + + +async def start(f: Callable[[list[str], mapstreamer.Datum], AsyncIterator[Message]]): + sock_file = "/tmp/var/run/numaflow/mapstream.sock" + server_info_file = "/tmp/var/run/numaflow/mapper-server-info" + server = mapstreamer.MapStreamAsyncServer(sock_file, server_info_file) + + # Register loop-level signal handlers to request graceful shutdown + 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 if __name__ == "__main__": - mapstreamer.MapStreamAsyncServer( - async_handler, - sock_file="/tmp/var/run/numaflow/mapstream.sock", - server_info_file="/tmp/var/run/numaflow/mapper-server-info", - ).run() + asyncio.run(start(async_handler)) diff --git a/packages/pynumaflow-lite/tests/examples/mapstream_cat_class.py b/packages/pynumaflow-lite/tests/examples/mapstream_cat_class.py index 8802b565..c90ffc75 100644 --- a/packages/pynumaflow-lite/tests/examples/mapstream_cat_class.py +++ b/packages/pynumaflow-lite/tests/examples/mapstream_cat_class.py @@ -1,23 +1,42 @@ -from collections.abc import AsyncIterable +import asyncio +import signal +from collections.abc import AsyncIterator, Callable from pynumaflow_lite import mapstreamer from pynumaflow_lite.mapstreamer import Message -class SimpleStreamCat: - async def handler(self, datum: mapstreamer.Datum) -> AsyncIterable[Message]: +class SimpleStreamCat(mapstreamer.MapStreamer): + async def handler(self, keys: list[str], datum: mapstreamer.Datum) -> AsyncIterator[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=datum.keys) + yield Message(s.encode(), keys) + + +async def start(f: Callable[[list[str], mapstreamer.Datum], AsyncIterator[Message]]): + sock_file = "/tmp/var/run/numaflow/mapstream.sock" + server_info_file = "/tmp/var/run/numaflow/mapper-server-info" + server = mapstreamer.MapStreamAsyncServer(sock_file, server_info_file) + + # 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 if __name__ == "__main__": - map_streamer_obj = SimpleStreamCat() - mapstreamer.MapStreamAsyncServer( - map_streamer_obj.handler, - sock_file="/tmp/var/run/numaflow/mapstream.sock", - server_info_file="/tmp/var/run/numaflow/mapper-server-info", - ).run() + async_handler = SimpleStreamCat() + asyncio.run(start(async_handler)) diff --git a/packages/pynumaflow-lite/tests/examples/sink_log_class.py b/packages/pynumaflow-lite/tests/examples/sink_log_class.py index 4e9a5e32..dda93180 100644 --- a/packages/pynumaflow-lite/tests/examples/sink_log_class.py +++ b/packages/pynumaflow-lite/tests/examples/sink_log_class.py @@ -1,21 +1,22 @@ import logging -from collections.abc import AsyncIterator +from collections.abc import AsyncIterable from pynumaflow_lite import sinker +from pynumaflow_lite.sinker import Sinker # Configure logging logging.basicConfig(level=logging.INFO) _LOGGER = logging.getLogger(__name__) -class SimpleLogSink: +class SimpleLogSink(Sinker): """ Simple log sink that logs each message and returns success responses. This is the class-based approach matching the user's example. Also demonstrates reading metadata (read-only for sink). """ - async def handler(self, datums: AsyncIterator[sinker.Datum]) -> list[sinker.Response]: + async def handler(self, datums: AsyncIterable[sinker.Datum]) -> list[sinker.Response]: responses = [] async for msg in datums: _LOGGER.info("User Defined Sink %s", msg.value.decode("utf-8")) @@ -39,9 +40,8 @@ async def handler(self, datums: AsyncIterator[sinker.Datum]) -> list[sinker.Resp if __name__ == "__main__": - sinker_obj = SimpleLogSink() sinker.SinkAsyncServer( - sinker_obj.handler, + SimpleLogSink(), sock_file="/tmp/var/run/numaflow/sink.sock", server_info_file="/tmp/var/run/numaflow/sinker-server-info", ).run() diff --git a/packages/pynumaflow-lite/tests/test_map_unit.py b/packages/pynumaflow-lite/tests/test_map_unit.py deleted file mode 100644 index 1701152c..00000000 --- a/packages/pynumaflow-lite/tests/test_map_unit.py +++ /dev/null @@ -1,92 +0,0 @@ -import asyncio -import os -import subprocess -import uuid -from contextlib import suppress -from pathlib import Path - -import pytest - -from pynumaflow_lite import mapper - -CARGO_ROOT = Path(__file__).resolve().parent.parent - - -async def _unit_handler(datum: mapper.Datum) -> list[mapper.Message]: - assert datum.value == b"value" - assert datum.user_metadata["custom"]["version"] == b"1.0.0" - return [mapper.Message(datum.value, keys=datum.keys)] - - -def test_map_handler_can_be_unit_tested_without_server(): - datum = mapper.Datum( - keys=["key"], - value=b"value", - user_metadata={"custom": {"version": b"1.0.0"}}, - ) - messages = asyncio.run(_unit_handler(datum)) - - assert messages == [mapper.Message(b"value", keys=["key"])] - - -def test_to_drop_message(): - msg = mapper.Message.to_drop() - - assert msg.value == b"" - assert msg.tags is not None - assert len(msg.tags) == 1 - - -async def _run_map_client(sock_path: Path) -> subprocess.CompletedProcess[str]: - return await asyncio.to_thread( - subprocess.run, - ["cargo", "run", "--quiet", "--bin", "test_map", "--", str(sock_path)], - cwd=CARGO_ROOT, - capture_output=True, - text=True, - timeout=60, - check=False, - ) - - -async def _exercise_server(tmp_path: Path, handler) -> None: - path_id = f"{tmp_path.name[:12]}-{uuid.uuid4().hex[:12]}" - sock_path = Path(f"/tmp/pnl-{path_id}.sock") - server_info_path = Path(f"/tmp/pnl-{path_id}.info") - server = mapper.MapAsyncServer( - handler, - sock_file=str(sock_path), - server_info_file=str(server_info_path), - ) - - try: - async with server: - await _run_map_client(sock_path) - finally: - for path in (sock_path, server_info_path): - with suppress(FileNotFoundError): - os.unlink(path) - - -def test_map_server_propagates_handler_exception(tmp_path: Path): - async def handler(datum: mapper.Datum) -> list[mapper.Message]: - raise RuntimeError("map exploded") - - with pytest.raises(RuntimeError, match="map exploded"): - asyncio.run(_exercise_server(tmp_path, handler)) - - -def test_map_server_rejects_non_list_response(tmp_path: Path): - async def handler(datum: mapper.Datum) -> str: - return "not messages" - - with pytest.raises(TypeError, match=r"map handler must return list\[Message\]"): - asyncio.run(_exercise_server(tmp_path, handler)) - - -def test_map_server_rejects_sync_handler(tmp_path: Path): - def handler(datum: mapper.Datum) -> list[mapper.Message]: - return [] - - with pytest.raises(TypeError, match="map handler must be an async function"): - asyncio.run(_exercise_server(tmp_path, handler)) From e28872f07cd9eef45ed1e28a9bec20975d3f4bf0 Mon Sep 17 00:00:00 2001 From: Vigith Maurice Date: Thu, 10 Sep 2026 18:22:38 -0700 Subject: [PATCH 2/2] fix(pynumaflow-lite): remove redundant borrows in println Signed-off-by: Vigith Maurice --- packages/pynumaflow-lite/src/sink/server.rs | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/packages/pynumaflow-lite/src/sink/server.rs b/packages/pynumaflow-lite/src/sink/server.rs index 671ac02f..f588bdf7 100644 --- a/packages/pynumaflow-lite/src/sink/server.rs +++ b/packages/pynumaflow-lite/src/sink/server.rs @@ -125,7 +125,7 @@ pub(super) async fn start( let py_asyncio_loop_handle = tokio::task::spawn_blocking({ println!( "Starting Sink UDF. socket={}, server_info={}", - &sock_file, &info_file + sock_file, info_file ); move || crate::pyrs::run_asyncio(tx) });