diff --git a/.github/workflows/pynumaflow-lite.yml b/.github/workflows/pynumaflow-lite.yml index 2d239831..feb79796 100644 --- a/.github/workflows/pynumaflow-lite.yml +++ b/.github/workflows/pynumaflow-lite.yml @@ -30,6 +30,7 @@ env: RUSTFLAGS: -C debuginfo=0 RUST_BACKTRACE: 1 PYTHONUTF8: 1 + RUST_TOOLCHAIN: "1.97" jobs: lint: @@ -55,8 +56,8 @@ jobs: - name: Set up Rust run: | - rustup toolchain install stable --profile minimal --component clippy --component rustfmt - rustup default stable + rustup toolchain install "${RUST_TOOLCHAIN}" --profile minimal --component clippy --component rustfmt + rustup default "${RUST_TOOLCHAIN}" - name: Install dependencies run: uv sync --locked --group dev @@ -100,8 +101,8 @@ jobs: - name: Set up Rust run: | - rustup toolchain install stable --profile minimal - rustup default stable + rustup toolchain install "${RUST_TOOLCHAIN}" --profile minimal + rustup default "${RUST_TOOLCHAIN}" - name: Run tests for Python versions shell: bash @@ -173,8 +174,8 @@ jobs: - name: Set up Rust run: | - rustup toolchain install stable --profile minimal - rustup default stable + rustup toolchain install "${RUST_TOOLCHAIN}" --profile minimal + rustup default "${RUST_TOOLCHAIN}" - name: Verify tag matches package version if: github.event_name == 'push' && startsWith(github.ref, 'refs/tags/pynumaflow-lite-v') diff --git a/packages/pynumaflow-lite/manifests/batchmap/batchmap_cat.py b/packages/pynumaflow-lite/manifests/batchmap/batchmap_cat.py index fed9eedf..9dff944a 100644 --- a/packages/pynumaflow-lite/manifests/batchmap/batchmap_cat.py +++ b/packages/pynumaflow-lite/manifests/batchmap/batchmap_cat.py @@ -1,46 +1,20 @@ -import asyncio -import signal -from collections.abc import AsyncIterable, Awaitable, Callable +from collections.abc import AsyncIterator from pynumaflow_lite import batchmapper -from pynumaflow_lite.batchmapper import Message +from pynumaflow_lite.batchmapper import BatchResponse, Datum, Message -class SimpleBatchCat(batchmapper.BatchMapper): - async def handler(self, batch: AsyncIterable[batchmapper.Datum]) -> batchmapper.BatchResponses: - responses = batchmapper.BatchResponses() - async for d in batch: - resp = batchmapper.BatchResponse(d.id) - if d.value == b"bad world": - resp.append(Message.message_to_drop()) - continue - - resp.append(Message(d.value, d.keys)) - responses.append(resp) - return responses - - -async def start( - f: Callable[[AsyncIterable[batchmapper.Datum]], Awaitable[batchmapper.BatchResponses]], -): - server = batchmapper.BatchMapAsyncServer() - - # Register loop-level signal handlers so we control shutdown and avoid asyncio.run - loop = asyncio.get_running_loop() - try: - loop.add_signal_handler(signal.SIGINT, lambda: server.stop()) - loop.add_signal_handler(signal.SIGTERM, lambda: server.stop()) - except (NotImplementedError, RuntimeError): - pass - - try: - await server.start(f) - print("Shutting down gracefully...") - except asyncio.CancelledError: - server.stop() - return +class SimpleBatchCat: + async def handler(self, batch: AsyncIterator[Datum]) -> list[BatchResponse]: + return [ + BatchResponse( + d.id, + [Message.to_drop()] if d.value == b"bad world" else [Message(d.value, keys=d.keys)], + ) + async for d in batch + ] if __name__ == "__main__": - async_handler = SimpleBatchCat() - asyncio.run(start(async_handler)) + batch_mapper_obj = SimpleBatchCat() + batchmapper.BatchMapAsyncServer(batch_mapper_obj.handler).run() diff --git a/packages/pynumaflow-lite/manifests/map/map_cat.py b/packages/pynumaflow-lite/manifests/map/map_cat.py index f877f212..8edc9cc5 100644 --- a/packages/pynumaflow-lite/manifests/map/map_cat.py +++ b/packages/pynumaflow-lite/manifests/map/map_cat.py @@ -1,48 +1,13 @@ -import asyncio -import signal -from collections.abc import Awaitable, Callable - from pynumaflow_lite import mapper -class SimpleCat(mapper.Mapper): - async def handler(self, keys: list[str], payload: mapper.Datum) -> mapper.Messages: - - messages = mapper.Messages() - - if payload.value == b"bad world": - messages.append(mapper.Message.message_to_drop()) - else: - messages.append(mapper.Message(payload.value, keys)) - - return messages - - -async def start(f: Callable[[list[str], mapper.Datum], Awaitable[mapper.Messages]]): - server = mapper.MapAsyncServer() - - # Register loop-level signal handlers so we control shutdown and avoid asyncio.run - # converting it into KeyboardInterrupt/CancelledError traces. - loop = asyncio.get_running_loop() - loop.set_debug(True) - print("Registering signal handlers", loop) - try: - loop.add_signal_handler(signal.SIGINT, lambda: server.stop()) - loop.add_signal_handler(signal.SIGTERM, lambda: server.stop()) - except (NotImplementedError, RuntimeError): - print("Failed to register signal handlers") - # add_signal_handler may not be available on some platforms/contexts; fallback below. - pass - - try: - await server.start(f) - print("Shutting down gracefully...") - except asyncio.CancelledError: - # Fallback in case the task was cancelled by the runner - server.stop() - return +class SimpleCat: + async def handler(self, datum: mapper.Datum) -> list[mapper.Message]: + if datum.value == b"bad world": + return [mapper.Message.to_drop()] + return [mapper.Message(datum.value, keys=datum.keys)] if __name__ == "__main__": - async_handler = SimpleCat() - asyncio.run(start(async_handler)) + mapper_obj = SimpleCat() + mapper.MapAsyncServer(mapper_obj.handler).run() diff --git a/packages/pynumaflow-lite/manifests/mapstream/mapstream_cat.py b/packages/pynumaflow-lite/manifests/mapstream/mapstream_cat.py index 9cfe0beb..84933c98 100644 --- a/packages/pynumaflow-lite/manifests/mapstream/mapstream_cat.py +++ b/packages/pynumaflow-lite/manifests/mapstream/mapstream_cat.py @@ -1,41 +1,19 @@ -import asyncio -import signal -from collections.abc import AsyncIterator, Callable +from collections.abc import AsyncIterable from pynumaflow_lite import mapstreamer from pynumaflow_lite.mapstreamer import Message -class SimpleStreamCat(mapstreamer.MapStreamer): - async def handler(self, keys: list[str], datum: mapstreamer.Datum) -> AsyncIterator[Message]: +class SimpleStreamCat: + async def handler(self, datum: mapstreamer.Datum) -> AsyncIterable[Message]: parts = datum.value.decode("utf-8").split(",") if not parts: yield Message.to_drop() return for s in parts: - yield Message(s.encode(), keys) - - -async def start(f: Callable[[list[str], mapstreamer.Datum], AsyncIterator[Message]]): - # Use default socket/info file locations; no explicit sock file passed - server = mapstreamer.MapStreamAsyncServer() - - # Register loop-level signal handlers so we control shutdown and avoid asyncio.run noise. - loop = asyncio.get_running_loop() - try: - loop.add_signal_handler(signal.SIGINT, lambda: server.stop()) - loop.add_signal_handler(signal.SIGTERM, lambda: server.stop()) - except (NotImplementedError, RuntimeError): - pass - - try: - await server.start(f) - print("Shutting down gracefully...") - except asyncio.CancelledError: - server.stop() - return + yield Message(s.encode(), keys=datum.keys) if __name__ == "__main__": - async_handler = SimpleStreamCat() - asyncio.run(start(async_handler)) + map_streamer_obj = SimpleStreamCat() + mapstreamer.MapStreamAsyncServer(map_streamer_obj.handler).run() diff --git a/packages/pynumaflow-lite/manifests/sideinput/sideinput_example.py b/packages/pynumaflow-lite/manifests/sideinput/sideinput_example.py index 5361b57f..8fb790a9 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(mapper.Mapper): +class SideInputHandler: """ A Mapper that reads from side input files and includes the value in its output. """ @@ -51,13 +51,11 @@ class SideInputHandler(mapper.Mapper): # Side input file that we are watching watched_file = "myticker" - async def handler(self, keys: list[str], datum: mapper.Datum) -> mapper.Messages: + async def handler(self, datum: mapper.Datum) -> list[mapper.Message]: with self.data_value_lock: current_value = self.data_value - messages = mapper.Messages() - messages.append(mapper.Message(str.encode(current_value))) - return messages + return [mapper.Message(str.encode(current_value))] def file_watcher(self): """ @@ -103,27 +101,18 @@ async def start_sideinput(): server.stop() -async def start_mapper(): +def start_mapper(): """Start the Mapper server that reads from side inputs.""" - server = mapper.MapAsyncServer() - handler = SideInputHandler() + mapper_obj = SideInputHandler() # Initialize the data value from the side input file - handler.init_data_value() + mapper_obj.init_data_value() # Start the file watcher in a background thread - watcher_thread = Thread(target=handler.file_watcher, daemon=True) + watcher_thread = Thread(target=mapper_obj.file_watcher, daemon=True) watcher_thread.start() - loop = asyncio.get_running_loop() - loop.add_signal_handler(signal.SIGINT, lambda: server.stop()) - loop.add_signal_handler(signal.SIGTERM, lambda: server.stop()) - - try: - await server.start(handler) - print("Mapper server shutting down gracefully...") - except asyncio.CancelledError: - server.stop() + mapper.MapAsyncServer(mapper_obj.handler).run() if __name__ == "__main__": @@ -132,7 +121,7 @@ async def start_mapper(): if is_mapper: print("Starting as Mapper (reading side inputs)...") - asyncio.run(start_mapper()) + start_mapper() else: print("Starting as SideInput retriever...") asyncio.run(start_sideinput()) diff --git a/packages/pynumaflow-lite/manifests/sink/sink_log.py b/packages/pynumaflow-lite/manifests/sink/sink_log.py index 1b1c656c..1a94aeff 100644 --- a/packages/pynumaflow-lite/manifests/sink/sink_log.py +++ b/packages/pynumaflow-lite/manifests/sink/sink_log.py @@ -1,20 +1,19 @@ import logging -from collections.abc import AsyncIterable +from collections.abc import AsyncIterator from pynumaflow_lite import sinker -from pynumaflow_lite.sinker import Sinker # Configure logging logging.basicConfig(level=logging.INFO) _LOGGER = logging.getLogger(__name__) -class SimpleLogSink(Sinker): +class SimpleLogSink: """ Simple log sink that logs each message and returns success responses. """ - async def handler(self, datums: AsyncIterable[sinker.Datum]) -> list[sinker.Response]: + async def handler(self, datums: AsyncIterator[sinker.Datum]) -> list[sinker.Response]: responses = [] async for msg in datums: _LOGGER.info("User Defined Sink: %s", msg.value.decode("utf-8")) @@ -25,4 +24,5 @@ async def handler(self, datums: AsyncIterable[sinker.Datum]) -> list[sinker.Resp if __name__ == "__main__": - sinker.SinkAsyncServer(SimpleLogSink()).run() + sinker_obj = SimpleLogSink() + sinker.SinkAsyncServer(sinker_obj.handler).run() diff --git a/packages/pynumaflow-lite/pynumaflow_lite/__init__.py b/packages/pynumaflow-lite/pynumaflow_lite/__init__.py index d1bfeb9f..30dab9f0 100644 --- a/packages/pynumaflow-lite/pynumaflow_lite/__init__.py +++ b/packages/pynumaflow-lite/pynumaflow_lite/__init__.py @@ -61,29 +61,27 @@ except Exception: # pragma: no cover sideinputer = None -# Surface the Python Mapper, BatchMapper, MapStreamer, Reducer, SessionReducer, ReduceStreamer, Accumulator, Sinker, -# Sourcer, SourceTransformer, and SideInput classes under the extension submodules for convenient access +# Surface the Python async servers and data-type classes under the extension submodules for convenient access from ._accumulator_dtypes import Accumulator -from ._batchmapper_dtypes import BatchMapper -from ._map_dtypes import Mapper -from ._mapstream_dtypes import MapStreamer +from ._batchmap_server import BatchMapAsyncServer +from ._map_server import MapAsyncServer +from ._mapstream_server import MapStreamAsyncServer from ._reduce_dtypes import Reducer from ._reducestreamer_dtypes import ReduceStreamer from ._session_reduce_dtypes import SessionReducer from ._sideinput_dtypes import SideInput -from ._sink_dtypes import Sinker from ._sink_server import SinkAsyncServer from ._source_dtypes import Sourcer from ._sourcetransformer_dtypes import SourceTransformer if mapper is not None: - mapper.Mapper = Mapper + mapper.MapAsyncServer = MapAsyncServer if batchmapper is not None: - batchmapper.BatchMapper = BatchMapper + batchmapper.BatchMapAsyncServer = BatchMapAsyncServer if mapstreamer is not None: - mapstreamer.MapStreamer = MapStreamer + mapstreamer.MapStreamAsyncServer = MapStreamAsyncServer if reducer is not None: reducer.Reducer = Reducer @@ -98,7 +96,6 @@ accumulator.Accumulator = Accumulator if sinker is not None: - sinker.Sinker = Sinker sinker.SinkAsyncServer = SinkAsyncServer if sourcer is not None: diff --git a/packages/pynumaflow-lite/pynumaflow_lite/_batchmap_server.py b/packages/pynumaflow-lite/pynumaflow_lite/_batchmap_server.py new file mode 100644 index 00000000..b39f5164 --- /dev/null +++ b/packages/pynumaflow-lite/pynumaflow_lite/_batchmap_server.py @@ -0,0 +1,91 @@ +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 deleted file mode 100644 index 557c8afe..00000000 --- a/packages/pynumaflow-lite/pynumaflow_lite/_batchmapper_dtypes.py +++ /dev/null @@ -1,21 +0,0 @@ -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 deleted file mode 100644 index 13733ed7..00000000 --- a/packages/pynumaflow-lite/pynumaflow_lite/_map_dtypes.py +++ /dev/null @@ -1,24 +0,0 @@ -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 new file mode 100644 index 00000000..4b3cb68e --- /dev/null +++ b/packages/pynumaflow-lite/pynumaflow_lite/_map_server.py @@ -0,0 +1,88 @@ +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 deleted file mode 100644 index 32c2b9cb..00000000 --- a/packages/pynumaflow-lite/pynumaflow_lite/_mapstream_dtypes.py +++ /dev/null @@ -1,22 +0,0 @@ -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 new file mode 100644 index 00000000..4349878d --- /dev/null +++ b/packages/pynumaflow-lite/pynumaflow_lite/_mapstream_server.py @@ -0,0 +1,88 @@ +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 deleted file mode 100644 index 840dd65f..00000000 --- a/packages/pynumaflow-lite/pynumaflow_lite/_sink_dtypes.py +++ /dev/null @@ -1,21 +0,0 @@ -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 4dd29892..5a4ef95d 100644 --- a/packages/pynumaflow-lite/pynumaflow_lite/_sink_server.py +++ b/packages/pynumaflow-lite/pynumaflow_lite/_sink_server.py @@ -2,16 +2,22 @@ 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: Any, + handler: Callable[ + [AsyncIterator[Datum]], + Awaitable[list[Response]], + ], *, sock_file: str | None = None, server_info_file: str | None = None, @@ -19,9 +25,16 @@ 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: - await self._core.start(self._handler) + if self._serving: + raise RuntimeError("sink 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() diff --git a/packages/pynumaflow-lite/pynumaflow_lite/batchmapper.pyi b/packages/pynumaflow-lite/pynumaflow_lite/batchmapper.pyi index 277af5e3..71cbcf9c 100644 --- a/packages/pynumaflow-lite/pynumaflow_lite/batchmapper.pyi +++ b/packages/pynumaflow-lite/pynumaflow_lite/batchmapper.pyi @@ -2,6 +2,7 @@ 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 @@ -15,48 +16,77 @@ class Message: tags: list[str] | None = ..., ) -> None: ... @staticmethod - def message_to_drop() -> Message: ... + def to_drop() -> Message: ... + def __repr__(self) -> str: ... + def __eq__(self, other: object) -> bool: ... class Datum: keys: list[str] value: bytes watermark: _dt.datetime - eventtime: _dt.datetime + event_time: _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) -> None: ... - @staticmethod - def from_id(id: str) -> BatchResponse: ... + def __init__(self, id: str, messages: list[Message] | None = ...) -> None: ... def append(self, message: Message) -> None: ... + def __len__(self) -> int: ... + def __repr__(self) -> str: ... + def __eq__(self, other: object) -> bool: ... -class BatchResponses: - def __init__(self) -> None: ... - def append(self, response: BatchResponse) -> None: ... +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 BatchMapAsyncServer: def __init__( self, + handler: Callable[[AsyncIterator[Datum]], Awaitable[list[BatchResponse]]], + *, sock_file: str | None = ..., - info_file: str | None = ..., + server_info_file: str | None = ..., ) -> None: ... - def start(self, py_func: Callable[[AsyncIterator[Datum]], Awaitable[BatchResponses]]) -> Awaitable[None]: ... + def run(self) -> None: ... + async def serve(self) -> None: ... def stop(self) -> None: ... - -class BatchMapper: - async def handler(self, batch: AsyncIterator[Datum]) -> BatchResponses: ... + 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: ... __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 8255e6d6..b5c432a6 100644 --- a/packages/pynumaflow-lite/pynumaflow_lite/mapper.pyi +++ b/packages/pynumaflow-lite/pynumaflow_lite/mapper.pyi @@ -2,114 +2,84 @@ from __future__ import annotations import datetime as _dt from collections.abc import Awaitable, Callable - -# 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: ... +from types import TracebackType class Message: keys: list[str] | None value: bytes tags: list[str] | None - user_metadata: UserMetadata | None + user_metadata: dict[str, dict[str, bytes]] | None def __init__( self, value: bytes, keys: list[str] | None = ..., tags: list[str] | None = ..., - user_metadata: UserMetadata | None = ..., + user_metadata: dict[str, dict[str, bytes]] | None = ..., ) -> None: ... @staticmethod - def message_to_drop() -> Message: ... + def to_drop() -> Message: ... + def __repr__(self) -> str: ... + def __eq__(self, other: object) -> bool: ... class Datum: - # Read-only attributes provided by the extension keys: list[str] value: bytes watermark: _dt.datetime - eventtime: _dt.datetime + event_time: _dt.datetime headers: dict[str, str] - user_metadata: UserMetadata - system_metadata: SystemMetadata + user_metadata: dict[str, dict[str, bytes]] + system_metadata: dict[str, dict[str, bytes]] + 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: +class _MapAsyncServer: def __init__( self, sock_file: str | None = ..., - info_file: str | None = ..., + server_info_file: str | None = ..., ) -> None: ... - def start(self, py_func: Callable[[list[str], Datum], Awaitable[Messages]]) -> Awaitable[None]: ... + def start( + self, + handler: Callable[[Datum], Awaitable[list[Message]]], + ) -> Awaitable[None]: ... + def wait_ready(self, timeout: float = ...) -> Awaitable[None]: ... def stop(self) -> None: ... -# Simple utility function exposed by the extension +class MapAsyncServer: + def __init__( + self, + handler: Callable[[Datum], Awaitable[list[Message]]], + *, + sock_file: str | None = ..., + server_info_file: str | None = ..., + ) -> None: ... + def run(self) -> None: ... + async def serve(self) -> 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: ... __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 ec9209af..d6660aa3 100644 --- a/packages/pynumaflow-lite/pynumaflow_lite/mapstreamer.pyi +++ b/packages/pynumaflow-lite/pynumaflow_lite/mapstreamer.pyi @@ -1,7 +1,8 @@ from __future__ import annotations import datetime as _dt -from collections.abc import AsyncIterator, Awaitable, Callable +from collections.abc import AsyncIterable, Awaitable, Callable +from types import TracebackType class Message: keys: list[str] | None @@ -15,35 +16,64 @@ 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 - eventtime: _dt.datetime + event_time: _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: +class _MapStreamAsyncServer: def __init__( self, sock_file: str | None = ..., - info_file: str | None = ..., + server_info_file: str | None = ..., ) -> None: ... - def start(self, py_func: Callable[[list[str], Datum], AsyncIterator[Message]]) -> Awaitable[None]: ... + def start( + self, + handler: Callable[[Datum], AsyncIterable[Message]], + ) -> Awaitable[None]: ... + def wait_ready(self, timeout: float = ...) -> Awaitable[None]: ... def stop(self) -> None: ... -class MapStreamer: - async def handler(self, keys: list[str], datum: Datum) -> AsyncIterator[Message]: ... +class MapStreamAsyncServer: + def __init__( + self, + handler: Callable[[Datum], AsyncIterable[Message]], + *, + sock_file: str | None = ..., + server_info_file: str | None = ..., + ) -> None: ... + def run(self) -> None: ... + async def serve(self) -> 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: ... __all__ = [ "Datum", "MapStreamAsyncServer", - "MapStreamer", "Message", ] diff --git a/packages/pynumaflow-lite/pynumaflow_lite/sinker.pyi b/packages/pynumaflow-lite/pynumaflow_lite/sinker.pyi index 5278f341..5be350e0 100644 --- a/packages/pynumaflow-lite/pynumaflow_lite/sinker.pyi +++ b/packages/pynumaflow-lite/pynumaflow_lite/sinker.pyi @@ -60,26 +60,23 @@ 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: _SinkHandler) -> Awaitable[None]: ... + def start( + self, + handler: Callable[[AsyncIterator[Datum]], Awaitable[list[Response]]], + ) -> 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: _SinkHandler | Sinker, + handler: Callable[[AsyncIterator[Datum]], Awaitable[list[Response]]], *, sock_file: str | None = ..., server_info_file: str | None = ..., @@ -101,5 +98,4 @@ __all__ = [ "Message", "Response", "SinkAsyncServer", - "Sinker", ] diff --git a/packages/pynumaflow-lite/src/batchmap/mod.rs b/packages/pynumaflow-lite/src/batchmap/mod.rs index fc43433b..e127110a 100644 --- a/packages/pynumaflow-lite/src/batchmap/mod.rs +++ b/packages/pynumaflow-lite/src/batchmap/mod.rs @@ -1,4 +1,5 @@ use std::collections::HashMap; +use std::time::Duration; use numaflow::batchmap; @@ -14,15 +15,18 @@ 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)] -#[derive(Clone, Default, Debug)] +#[pyclass(module = "pynumaflow_lite.batchmapper", from_py_object, eq)] +#[derive(Clone, Default, Debug, PartialEq)] 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>, } @@ -37,15 +41,28 @@ impl Message { } /// Drop a [Message], do not forward to the next vertex. - #[pyo3(signature = ())] #[staticmethod] - fn message_to_drop() -> Self { + #[pyo3(signature = () -> "Message")] + fn 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 { @@ -72,7 +89,7 @@ pub struct Datum { pub watermark: DateTime, /// Time of the element as seen at source or aligned after a reduce operation. #[pyo3(get)] - pub eventtime: DateTime, + pub event_time: DateTime, /// ID is the unique id of the message #[pyo3(get)] pub id: String, @@ -81,112 +98,104 @@ 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: Vec, - value: Vec, - watermark: DateTime, - eventtime: DateTime, - id: String, - headers: HashMap, + keys: Option>, + value: Option>, + id: Option, + event_time: Option>, + watermark: Option>, + headers: Option>, ) -> Self { Self { - keys, - value, - watermark, - eventtime, - id, - headers, + 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(), } } fn __repr__(&self) -> String { format!( - "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={:?})", + "Datum(keys={:?}, value={}, watermark={}, event_time={}, id={:?}, headers={:?})", self.keys, - String::from_utf8_lossy(&self.value), + crate::map::bytes_literal(&self.value), self.watermark, - self.eventtime, + self.event_time, self.id, self.headers ) } + + fn __str__(&self) -> String { + self.__repr__() + } } impl From for Datum { fn from(value: batchmap::Datum) -> Self { - Datum::new( - value.keys, - value.value, - value.watermark, - value.event_time, - value.id, - value.headers, - ) + Self { + keys: value.keys, + value: value.value, + watermark: value.watermark, + event_time: value.event_time, + id: value.id, + headers: value.headers, + } } } /// BatchResponse mirrors numaflow::batchmap::BatchResponse for Python -#[pyclass(module = "pynumaflow_lite.batchmapper", from_py_object)] -#[derive(Clone, Debug)] +#[pyclass(module = "pynumaflow_lite.batchmapper", from_py_object, eq)] +#[derive(Clone, Debug, PartialEq)] pub struct BatchResponse { #[pyo3(get)] pub id: String, + #[pyo3(get)] pub messages: Vec, } #[pymethods] impl BatchResponse { #[new] - #[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 { + #[pyo3(signature = (id: "str", messages: "list[Message] | None"=None) -> "BatchResponse")] + fn new(id: String, messages: Option>) -> Self { Self { id, - messages: Vec::new(), + messages: messages.unwrap_or_default(), } } - #[pyo3(signature = (message))] + #[pyo3(signature = (message: "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, -} -#[pymethods] -impl BatchResponses { - #[new] - #[pyo3(signature = () -> "BatchResponses")] - fn new() -> Self { - Self { responses: vec![] } + fn __len__(&self) -> usize { + self.messages.len() } - /// Append a BatchResponse to the collection. - #[pyo3(signature = (response: "BatchResponse"))] - fn append(&mut self, response: BatchResponse) { - self.responses.push(response); + fn __repr__(&self) -> String { + let messages = self + .messages + .iter() + .map(|message| message.__repr__()) + .collect::>() + .join(", "); + format!("BatchResponse(id={:?}, messages=[{}])", self.id, messages) } } @@ -209,14 +218,6 @@ 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 } @@ -235,30 +236,33 @@ impl PyAsyncDatumStream { } /// Async Batch Map Server that can be started from Python code -#[pyclass(module = "pynumaflow_lite.batchmapper")] +#[pyclass(name = "_BatchMapAsyncServer", module = "pynumaflow_lite.batchmapper")] pub struct BatchMapAsyncServer { sock_file: String, - info_file: String, + server_info_file: String, shutdown_tx: Mutex>>, } #[pymethods] impl BatchMapAsyncServer { #[new] - #[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 { + #[pyo3(signature = ( + sock_file: "str | None"=None, + server_info_file: "str | None"=None, + ) -> "_BatchMapAsyncServer")] + fn new(sock_file: Option, server_info_file: Option) -> Self { Self { - sock_file, - info_file, + 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()), shutdown_tx: Mutex::new(None), } } - #[pyo3(signature = (py_func: "callable") -> "None")] - pub fn start<'a>(&self, py: Python<'a>, py_func: Py) -> PyResult> { + #[pyo3(signature = (handler: "callable") -> "None")] + pub fn start<'a>(&self, py: Python<'a>, handler: Py) -> PyResult> { let sock_file = self.sock_file.clone(); - let info_file = self.info_file.clone(); + let server_info_file = self.server_info_file.clone(); let (tx, rx) = tokio::sync::oneshot::channel::<()>(); { let mut guard = self.shutdown_tx.lock().unwrap(); @@ -266,11 +270,25 @@ impl BatchMapAsyncServer { } pyo3_async_runtimes::tokio::future_into_py(py, async move { - // 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"); + 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?; Ok(()) }) } @@ -289,9 +307,7 @@ 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 25c75210..dd799146 100644 --- a/packages/pynumaflow-lite/src/batchmap/server.rs +++ b/packages/pynumaflow-lite/src/batchmap/server.rs @@ -1,13 +1,34 @@ -// use crate::batchmap::Datum; use numaflow::batchmap; use numaflow::shared::ServerExtras; +use pyo3::exceptions::PyTypeError; use pyo3::prelude::*; -use std::sync::Arc; +use std::sync::{Arc, Mutex}; 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] @@ -29,28 +50,58 @@ impl batchmap::BatchMapper for PyBatchMapRunner { // When input ends, dropping tx closes the channel }); - // Call the Python coroutine: py_func(batch: AsyncIterable[Datum]) -> BatchResponses - let fut = Python::attach(|py| { + // Call the Python coroutine: py_func(batch: AsyncIterator[Datum]) -> list[BatchResponse] + let fut = match Python::attach(|py| -> PyResult<_> { 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,)).unwrap().into_bound(py); - pyo3_async_runtimes::into_future_with_locals(&locals, coro).unwrap() - }); - - let result = fut.await.unwrap(); + 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), + }; // Ensure forwarder completes let _ = forwarder.await; - let responses = Python::attach(|py| { - let x: crate::batchmap::BatchResponses = result.extract(py).unwrap(); - x - }); + 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), + }; responses - .responses .into_iter() .map(|resp| resp.into()) .collect::>() @@ -65,14 +116,34 @@ 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(move || crate::pyrs::run_asyncio(tx)); + 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 event_loop = rx.await.unwrap(); - let (sig_handle, combined_rx) = crate::pyrs::setup_sig_handler(shutdown_rx); + 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 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) @@ -91,7 +162,7 @@ pub(super) async fn start( } }); - println!("Numaflow Core (batch) has shutdown..."); + println!("Numaflow BatchMap has shutdown..."); // Wait for the blocking asyncio thread to finish. let _ = py_asyncio_loop_handle.await; @@ -102,5 +173,9 @@ 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 85c3843c..0232cf00 100644 --- a/packages/pynumaflow-lite/src/map/mod.rs +++ b/packages/pynumaflow-lite/src/map/mod.rs @@ -1,6 +1,9 @@ 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}; @@ -8,233 +11,89 @@ 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; -/// 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>>, -} - -#[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>>, +pub(crate) fn bytes_literal(value: &[u8]) -> String { + format!("b\"{}\"", String::from_utf8_lossy(value).escape_debug()) } -#[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 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(", ")) } -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); +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)); } - Self { data } + data.insert(group, kv); } + data } -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); - } +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)); } - 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.insert(group, kv); } + data } /// A message to be sent to the next vertex. -#[pyclass(module = "pynumaflow_lite.mapper", from_py_object)] -#[derive(Clone, Default, Debug)] +#[pyclass(module = "pynumaflow_lite.mapper", from_py_object, eq)] +#[derive(Clone, Default, Debug, PartialEq)] 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. - pub user_metadata: Option, + #[pyo3(get)] + pub user_metadata: Option>>>, } #[pymethods] impl Message { - /// Create a new [Message] with the given value, keys, tags, and user_metadata. + /// Create a new Message with the given value. Keys, tags, and user_metadata are optional. #[new] - #[pyo3(signature = (value: "bytes", keys: "list[str] | None"=None, tags: "list[str] | None"=None, user_metadata: "UserMetadata | None"=None) -> "Message" - )] + #[pyo3(signature = (value: "bytes", keys: "list[str] | None"=None, tags: "list[str] | None"=None, user_metadata: "dict[str, dict[str, bytes]] | None"=None) -> "Message")] fn new( value: Vec, keys: Option>, tags: Option>, - user_metadata: Option, + user_metadata: Option>>>, ) -> Self { Self { keys, @@ -244,10 +103,10 @@ impl Message { } } - /// Drop a [Message], do not forward to the next vertex. - #[pyo3(signature = ())] + /// A Message marked to be dropped, i.e. not forwarded to the next vertex. #[staticmethod] - fn message_to_drop() -> Self { + #[pyo3(signature = () -> "Message")] + fn to_drop() -> Self { Self { keys: None, value: vec![], @@ -255,6 +114,22 @@ 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 { @@ -263,12 +138,21 @@ impl From for map::Message { keys: value.keys, value: value.value, tags: value.tags, - user_metadata: value.user_metadata.map(|m| m.into()), + 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 + }), } } } -/// The incoming [MapRequest] accessible in Python function. +/// The incoming Datum passed to the map handler. It carries the event's keys, value, +/// event_time, watermark, headers, and the user/system metadata. #[pyclass(module = "pynumaflow_lite.mapper")] pub struct Datum { /// Set of keys in the (key, value) terminology of map/reduce paradigm. @@ -283,106 +167,162 @@ pub struct Datum { pub watermark: DateTime, /// Time of the element as seen at source or aligned after a reduce operation. #[pyo3(get)] - pub eventtime: DateTime, + pub event_time: DateTime, /// Headers for the message. #[pyo3(get)] pub headers: HashMap, /// User metadata for the message. #[pyo3(get)] - pub user_metadata: UserMetadata, + pub user_metadata: HashMap>>, /// System metadata for the message. #[pyo3(get)] - pub system_metadata: SystemMetadata, + pub system_metadata: HashMap>>, } +#[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: Vec, - value: Vec, - watermark: DateTime, - eventtime: DateTime, - headers: HashMap, - user_metadata: UserMetadata, - system_metadata: SystemMetadata, + keys: Option>, + value: Option>, + event_time: Option>, + watermark: Option>, + headers: Option>, + user_metadata: Option>>>, + system_metadata: Option>>>, ) -> Self { Self { - keys, - value, - watermark, - eventtime, - headers, - user_metadata, - system_metadata, + 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(), } } fn __repr__(&self) -> String { format!( - "Datum(keys={:?}, value={:?}, watermark={}, eventtime={}, headers={:?}, user_metadata={:?}, system_metadata={:?})", + "Datum(keys={:?}, value={}, watermark={}, event_time={}, headers={:?}, user_metadata={}, system_metadata={})", self.keys, - self.value, + bytes_literal(&self.value), self.watermark, - self.eventtime, + self.event_time, self.headers, - self.user_metadata, - self.system_metadata + metadata_literal(&self.user_metadata), + metadata_literal(&self.system_metadata) ) } fn __str__(&self) -> String { - 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 - ) + self.__repr__() } } impl From for Datum { fn from(value: map::MapRequest) -> Self { - Datum::new( - value.keys, - value.value, - value.watermark, - value.eventtime, - value.headers, - value.user_metadata.into(), - value.system_metadata.into(), - ) + 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; } } /// Async Map Server that can be started from Python code which will run the Python UDF function. -#[pyclass(module = "pynumaflow_lite.mapper")] +#[pyclass(name = "_MapAsyncServer", module = "pynumaflow_lite.mapper")] pub struct MapAsyncServer { sock_file: String, - info_file: String, + server_info_file: String, shutdown_tx: Mutex>>, } #[pymethods] impl MapAsyncServer { #[new] - #[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 { + #[pyo3(signature = ( + sock_file: "str | None"=None, + server_info_file: "str | None"=None, + ) -> "_MapAsyncServer")] + fn new(sock_file: Option, server_info_file: Option) -> Self { Self { - sock_file, - info_file, + 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()), shutdown_tx: Mutex::new(None), } } /// Start the server with the given Python function. - #[pyo3(signature = (py_func: "callable") -> "None")] - pub fn start<'a>(&self, py: Python<'a>, py_func: Py) -> PyResult> { + #[pyo3(signature = (handler: "callable") -> "None")] + pub fn start<'a>(&self, py: Python<'a>, handler: Py) -> PyResult> { let sock_file = self.sock_file.clone(); - let info_file = self.info_file.clone(); + let server_info_file = self.server_info_file.clone(); let (tx, rx) = tokio::sync::oneshot::channel::<()>(); { let mut guard = self.shutdown_tx.lock().unwrap(); @@ -390,9 +330,25 @@ impl MapAsyncServer { } pyo3_async_runtimes::tokio::future_into_py(py, async move { - crate::map::server::start(py_func, sock_file, info_file, rx) - .await - .expect("server failed to start"); + 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?; Ok(()) }) } @@ -409,9 +365,6 @@ 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 1e220994..98a7a194 100644 --- a/packages/pynumaflow-lite/src/map/server.rs +++ b/packages/pynumaflow-lite/src/map/server.rs @@ -1,40 +1,83 @@ -use crate::map::{Datum, Messages}; use numaflow::map; use numaflow::shared::ServerExtras; +use pyo3::exceptions::PyTypeError; use pyo3::prelude::*; -use std::sync::Arc; +use std::sync::{Arc, Mutex}; 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 = Python::attach(|py| { - let keys = input.keys.clone(); - let input: Datum = input.into(); - let py_func = self.py_func.clone(); - + let fut = match Python::attach(|py| -> PyResult<_> { let locals = pyo3_async_runtimes::TaskLocals::new(self.event_loop.bind(py).clone()); - - 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() + 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() } } @@ -46,14 +89,34 @@ 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(move || crate::pyrs::run_asyncio(tx)); + 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 event_loop = rx.await.unwrap(); - let (sig_handle, combined_rx) = crate::pyrs::setup_sig_handler(shutdown_rx); + 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 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) @@ -72,7 +135,7 @@ pub(super) async fn start( } }); - println!("Numaflow Core has shutdown..."); + println!("Numaflow Map has shutdown..."); // Wait for the blocking asyncio thread to finish. let _ = py_asyncio_loop_handle.await; @@ -83,5 +146,9 @@ 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 24e126f0..84288c57 100644 --- a/packages/pynumaflow-lite/src/mapstream/mod.rs +++ b/packages/pynumaflow-lite/src/mapstream/mod.rs @@ -2,6 +2,7 @@ use chrono::{DateTime, Utc}; use numaflow::mapstream; use std::collections::HashMap; use std::sync::Mutex; +use std::time::Duration; pub mod server; @@ -9,8 +10,7 @@ pub mod server; use pyo3::prelude::*; /// Streaming Datum mirrors MapStreamRequest for Python -#[pyclass(module = "pynumaflow_lite.mapstreamer", from_py_object)] -#[derive(Clone)] +#[pyclass(module = "pynumaflow_lite.mapstreamer")] pub struct Datum { /// Set of keys in the (key, value) terminology of the map/reduce paradigm. #[pyo3(get)] @@ -24,33 +24,76 @@ pub struct Datum { pub watermark: DateTime, /// Time of the element as seen at source or aligned after a reduce operation. #[pyo3(get)] - pub eventtime: DateTime, + pub event_time: DateTime, /// Headers associated with the message. #[pyo3(get)] pub headers: HashMap, } impl Datum { - pub(crate) fn new( + pub(crate) fn from_parts( keys: Vec, value: Vec, watermark: DateTime, - eventtime: DateTime, + event_time: DateTime, headers: HashMap, ) -> Self { Self { keys, value, watermark, - eventtime, + event_time, 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::new( + Self::from_parts( value.keys, value.value, value.watermark, @@ -61,14 +104,17 @@ impl From for Datum { } /// A message to be sent downstream from a streaming handler. -#[pyclass(module = "pynumaflow_lite.mapstreamer", from_py_object)] -#[derive(Clone, Default, Debug)] +#[pyclass(module = "pynumaflow_lite.mapstreamer", from_py_object, eq)] +#[derive(Clone, Default, Debug, PartialEq)] 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>, } @@ -81,10 +127,10 @@ impl Message { Self { keys, value, tags } } - /// Drop a Message, do not forward to the next vertex. - #[pyo3(signature = ())] + /// A Message marked to be dropped, i.e. not forwarded to the next vertex. #[staticmethod] - fn message_to_drop() -> Self { + #[pyo3(signature = () -> "Message")] + fn to_drop() -> Self { Self { keys: None, value: vec![], @@ -92,11 +138,17 @@ impl Message { } } - /// Convenience alias to match example usage: Message.to_drop() - #[pyo3(signature = ())] - #[staticmethod] - fn to_drop() -> Self { - Self::message_to_drop() + 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:?}")), + ) } } @@ -111,31 +163,34 @@ impl From for mapstream::Message { } /// Async MapStream Server that can be started from Python code which will run the Python UDF async generator. -#[pyclass(module = "pynumaflow_lite.mapstreamer")] +#[pyclass(name = "_MapStreamAsyncServer", module = "pynumaflow_lite.mapstreamer")] pub struct MapStreamAsyncServer { sock_file: String, - info_file: String, + server_info_file: String, shutdown_tx: Mutex>>, } #[pymethods] impl MapStreamAsyncServer { #[new] - #[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 { + #[pyo3(signature = ( + sock_file: "str | None"=None, + server_info_file: "str | None"=None, + ) -> "_MapStreamAsyncServer")] + fn new(sock_file: Option, server_info_file: Option) -> Self { Self { - sock_file, - info_file, + 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()), shutdown_tx: Mutex::new(None), } } /// Start the server with the given Python async generator function. - #[pyo3(signature = (py_func: "callable") -> "None")] - pub fn start<'a>(&self, py: Python<'a>, py_func: Py) -> PyResult> { + #[pyo3(signature = (handler: "callable") -> "None")] + pub fn start<'a>(&self, py: Python<'a>, handler: Py) -> PyResult> { let sock_file = self.sock_file.clone(); - let info_file = self.info_file.clone(); + let server_info_file = self.server_info_file.clone(); let (tx, rx) = tokio::sync::oneshot::channel::<()>(); { let mut guard = self.shutdown_tx.lock().unwrap(); @@ -143,9 +198,25 @@ impl MapStreamAsyncServer { } pyo3_async_runtimes::tokio::future_into_py(py, async move { - crate::mapstream::server::start(py_func, sock_file, info_file, rx) - .await - .expect("server failed to start"); + 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?; Ok(()) }) } diff --git a/packages/pynumaflow-lite/src/mapstream/server.rs b/packages/pynumaflow-lite/src/mapstream/server.rs index 7dbcaee1..2edde058 100644 --- a/packages/pynumaflow-lite/src/mapstream/server.rs +++ b/packages/pynumaflow-lite/src/mapstream/server.rs @@ -5,34 +5,65 @@ use crate::pyiterables::PyAsyncIterStream; use numaflow::mapstream; use numaflow::shared::ServerExtras; +use pyo3::exceptions::PyTypeError; use pyo3::prelude::*; -use std::sync::Arc; +use std::sync::{Arc, Mutex}; 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(keys, datum) -> AsyncIterator - let agen_obj = Python::attach(|py| { - let keys = input.keys.clone(); + // Call Python handler: handler(datum) -> AsyncIterable[Message] + let agen_obj = match Python::attach(|py| -> PyResult> { let datum: Datum = input.into(); let py_func = self.py_func.clone(); - 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) - }); + 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; + } + }; // Wrap the Python AsyncIterable in a Rust Stream that yields incrementally - let mut stream = PyAsyncIterStream::::new(agen_obj, self.event_loop.clone()) - .expect("failed to construct PyAsyncIterStream"); + let mut stream = + match PyAsyncIterStream::::new(agen_obj, self.event_loop.clone()) { + Ok(stream) => stream, + Err(error) => { + self.fail(error); + return; + } + }; // Forward each yielded message immediately to the sender while let Some(item) = stream.next().await { @@ -44,8 +75,7 @@ impl mapstream::MapStreamer for PyMapStreamRunner { } } Err(e) => { - // Non-stop errors are surfaced per-item; log and stop this stream. - eprintln!("Python async iteration error: {:?}", e); + self.fail(e); break; } } @@ -57,23 +87,43 @@ impl mapstream::MapStreamer for PyMapStreamRunner { pub(super) async fn start( py_func: Py, sock_file: String, - info_file: String, + server_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(move || crate::pyrs::run_asyncio(tx)); + 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 event_loop = rx.await.unwrap(); - let (sig_handle, combined_rx) = crate::pyrs::setup_sig_handler(shutdown_rx); + 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 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(info_file); + .with_server_info_file(server_info_file); let result = server .start_with_shutdown(combined_rx) @@ -87,7 +137,7 @@ pub(super) async fn start( } }); - println!("Numaflow Core (stream) has shutdown..."); + println!("Numaflow MapStream has shutdown..."); // Wait for the blocking asyncio thread to finish. let _ = py_asyncio_loop_handle.await; @@ -98,5 +148,9 @@ 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 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) }); diff --git a/packages/pynumaflow-lite/tests/_test_utils.py b/packages/pynumaflow-lite/tests/_test_utils.py index fced4816..3d272280 100644 --- a/packages/pynumaflow-lite/tests/_test_utils.py +++ b/packages/pynumaflow-lite/tests/_test_utils.py @@ -37,6 +37,16 @@ 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, @@ -91,6 +101,8 @@ 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 64a018e8..1558ca73 100644 --- a/packages/pynumaflow-lite/tests/examples/batchmap_cat.py +++ b/packages/pynumaflow-lite/tests/examples/batchmap_cat.py @@ -1,51 +1,24 @@ -import asyncio -import collections.abc -import signal -from collections.abc import Awaitable, Callable +from collections.abc import AsyncIterator from pynumaflow_lite import batchmapper +from pynumaflow_lite.batchmapper import BatchResponse, Datum, Message -async def async_handler( - batch: collections.abc.AsyncIterator[batchmapper.Datum], -) -> batchmapper.BatchResponses: - responses = batchmapper.BatchResponses() +async def async_handler(batch: AsyncIterator[Datum]) -> list[BatchResponse]: + responses = [] async for d in batch: - resp = batchmapper.BatchResponse.from_id(d.id) + resp = BatchResponse(d.id) if d.value == b"bad world": - resp.append(batchmapper.Message.message_to_drop()) - continue - - resp.append(batchmapper.Message(d.value, d.keys)) + resp.append(Message.to_drop()) + else: + resp.append(Message(d.value, keys=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__": - asyncio.run(start(async_handler)) + batchmapper.BatchMapAsyncServer( + async_handler, + sock_file="/tmp/var/run/numaflow/batchmap.sock", + server_info_file="/tmp/var/run/numaflow/mapper-server-info", + ).run() diff --git a/packages/pynumaflow-lite/tests/examples/batchmap_cat_class.py b/packages/pynumaflow-lite/tests/examples/batchmap_cat_class.py index f78cfdfe..e3dd6dd1 100644 --- a/packages/pynumaflow-lite/tests/examples/batchmap_cat_class.py +++ b/packages/pynumaflow-lite/tests/examples/batchmap_cat_class.py @@ -1,48 +1,24 @@ -import asyncio -import signal -from collections.abc import AsyncIterator, Awaitable, Callable +from collections.abc import AsyncIterator from pynumaflow_lite import batchmapper -from pynumaflow_lite.batchmapper import Message +from pynumaflow_lite.batchmapper import BatchResponse, Datum, Message -class SimpleBatchCat(batchmapper.BatchMapper): - async def handler(self, batch: 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 +class SimpleBatchCat: + async def handler(self, batch: AsyncIterator[Datum]) -> list[BatchResponse]: + return [ + BatchResponse( + d.id, + [Message.to_drop()] if d.value == b"bad world" else [Message(d.value, keys=d.keys)], + ) + async for d in batch + ] if __name__ == "__main__": - async_handler = SimpleBatchCat() - asyncio.run(start(async_handler)) + batch_mapper_obj = SimpleBatchCat() + batchmapper.BatchMapAsyncServer( + batch_mapper_obj.handler, + sock_file="/tmp/var/run/numaflow/batchmap.sock", + server_info_file="/tmp/var/run/numaflow/mapper-server-info", + ).run() diff --git a/packages/pynumaflow-lite/tests/examples/map_cat.py b/packages/pynumaflow-lite/tests/examples/map_cat.py index 9e04728b..f7947540 100644 --- a/packages/pynumaflow-lite/tests/examples/map_cat.py +++ b/packages/pynumaflow-lite/tests/examples/map_cat.py @@ -1,61 +1,34 @@ -import asyncio -import signal -from collections.abc import Awaitable, Callable - from pynumaflow_lite import mapper -async def async_handler(keys: list[str], payload: mapper.Datum) -> mapper.Messages: - messages = mapper.Messages() - +async def async_handler(datum: mapper.Datum) -> list[mapper.Message]: # Read system metadata (read-only) - 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 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[{group}][{key}] = {value}") - # 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) + # 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(): print(f" User[{group}][{key}] = {value}") - 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 + if datum.value == b"bad world": + return [mapper.Message.to_drop()] - try: - await server.start(f) - print("Shutting down gracefully...") - except asyncio.CancelledError: - server.stop() - return + 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)] if __name__ == "__main__": - asyncio.run(start(async_handler)) + mapper.MapAsyncServer( + async_handler, + sock_file="/tmp/var/run/numaflow/map.sock", + server_info_file="/tmp/var/run/numaflow/mapper-server-info", + ).run() diff --git a/packages/pynumaflow-lite/tests/examples/map_cat_class.py b/packages/pynumaflow-lite/tests/examples/map_cat_class.py index 079beaf2..b9b76ec6 100644 --- a/packages/pynumaflow-lite/tests/examples/map_cat_class.py +++ b/packages/pynumaflow-lite/tests/examples/map_cat_class.py @@ -1,70 +1,36 @@ -import asyncio -import signal -from collections.abc import Awaitable, Callable - from pynumaflow_lite import mapper -class SimpleCat(mapper.Mapper): - async def handler(self, keys: list[str], payload: mapper.Datum) -> mapper.Messages: - - messages = mapper.Messages() - +class SimpleCat: + async def handler(self, datum: mapper.Datum) -> list[mapper.Message]: # Read system metadata (read-only) - 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 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[{group}][{key}] = {value}") - # 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) + # 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(): print(f" User[{group}][{key}] = {value}") - 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 + if datum.value == b"bad world": + return [mapper.Message.to_drop()] - 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 + 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)] if __name__ == "__main__": - async_handler = SimpleCat() - asyncio.run(start(async_handler)) + 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() diff --git a/packages/pynumaflow-lite/tests/examples/mapstream_cat.py b/packages/pynumaflow-lite/tests/examples/mapstream_cat.py index 0fd0bcbb..d809d7e3 100644 --- a/packages/pynumaflow-lite/tests/examples/mapstream_cat.py +++ b/packages/pynumaflow-lite/tests/examples/mapstream_cat.py @@ -1,12 +1,10 @@ -import asyncio -import signal -from collections.abc import AsyncIterator, Callable +from collections.abc import AsyncIterable from pynumaflow_lite import mapstreamer from pynumaflow_lite.mapstreamer import Message -async def async_handler(keys: list[str], datum: mapstreamer.Datum) -> AsyncIterator[Message]: +async def async_handler(datum: mapstreamer.Datum) -> AsyncIterable[Message]: """ A handler that splits the input datum value into multiple strings by `,` separator and emits them as a stream. @@ -16,29 +14,12 @@ async def async_handler(keys: list[str], datum: mapstreamer.Datum) -> AsyncItera yield Message.to_drop() return for s in parts: - yield Message(s.encode(), keys) - - -async def start(f: Callable[[list[str], mapstreamer.Datum], AsyncIterator[Message]]): - 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 + yield Message(s.encode(), keys=datum.keys) if __name__ == "__main__": - asyncio.run(start(async_handler)) + mapstreamer.MapStreamAsyncServer( + async_handler, + sock_file="/tmp/var/run/numaflow/mapstream.sock", + server_info_file="/tmp/var/run/numaflow/mapper-server-info", + ).run() diff --git a/packages/pynumaflow-lite/tests/examples/mapstream_cat_class.py b/packages/pynumaflow-lite/tests/examples/mapstream_cat_class.py index c90ffc75..8802b565 100644 --- a/packages/pynumaflow-lite/tests/examples/mapstream_cat_class.py +++ b/packages/pynumaflow-lite/tests/examples/mapstream_cat_class.py @@ -1,42 +1,23 @@ -import asyncio -import signal -from collections.abc import AsyncIterator, Callable +from collections.abc import AsyncIterable from pynumaflow_lite import mapstreamer from pynumaflow_lite.mapstreamer import Message -class SimpleStreamCat(mapstreamer.MapStreamer): - async def handler(self, keys: list[str], datum: mapstreamer.Datum) -> AsyncIterator[Message]: +class SimpleStreamCat: + async def handler(self, datum: mapstreamer.Datum) -> AsyncIterable[Message]: parts = datum.value.decode("utf-8").split(",") if not parts: yield Message.to_drop() return for s in parts: - yield Message(s.encode(), keys) - - -async def start(f: Callable[[list[str], mapstreamer.Datum], AsyncIterator[Message]]): - 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 + yield Message(s.encode(), keys=datum.keys) if __name__ == "__main__": - async_handler = SimpleStreamCat() - asyncio.run(start(async_handler)) + map_streamer_obj = SimpleStreamCat() + mapstreamer.MapStreamAsyncServer( + map_streamer_obj.handler, + sock_file="/tmp/var/run/numaflow/mapstream.sock", + server_info_file="/tmp/var/run/numaflow/mapper-server-info", + ).run() diff --git a/packages/pynumaflow-lite/tests/examples/sink_log_class.py b/packages/pynumaflow-lite/tests/examples/sink_log_class.py index dda93180..4e9a5e32 100644 --- a/packages/pynumaflow-lite/tests/examples/sink_log_class.py +++ b/packages/pynumaflow-lite/tests/examples/sink_log_class.py @@ -1,22 +1,21 @@ import logging -from collections.abc import AsyncIterable +from collections.abc import AsyncIterator from pynumaflow_lite import sinker -from pynumaflow_lite.sinker import Sinker # Configure logging logging.basicConfig(level=logging.INFO) _LOGGER = logging.getLogger(__name__) -class SimpleLogSink(Sinker): +class SimpleLogSink: """ Simple log sink that logs each message and returns success responses. This is the class-based approach matching the user's example. Also demonstrates reading metadata (read-only for sink). """ - async def handler(self, datums: AsyncIterable[sinker.Datum]) -> list[sinker.Response]: + async def handler(self, datums: AsyncIterator[sinker.Datum]) -> list[sinker.Response]: responses = [] async for msg in datums: _LOGGER.info("User Defined Sink %s", msg.value.decode("utf-8")) @@ -40,8 +39,9 @@ async def handler(self, datums: AsyncIterable[sinker.Datum]) -> list[sinker.Resp if __name__ == "__main__": + sinker_obj = SimpleLogSink() sinker.SinkAsyncServer( - SimpleLogSink(), + sinker_obj.handler, 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 new file mode 100644 index 00000000..1701152c --- /dev/null +++ b/packages/pynumaflow-lite/tests/test_map_unit.py @@ -0,0 +1,92 @@ +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))