From 7de7071e1678b2cbe791783c74f2e44e0b1931e8 Mon Sep 17 00:00:00 2001 From: Sreekanth Date: Wed, 14 Oct 2026 06:54:14 +0530 Subject: [PATCH 1/2] pynumaflow-lite: Pythonic APIs for Reducer Signed-off-by: Sreekanth --- .../manifests/reduce/.dockerignore | 1 + .../manifests/reduce/Dockerfile | 74 +++--- .../manifests/reduce/pyproject.toml | 7 +- .../manifests/reduce/reduce_counter_class.py | 59 ++--- .../pynumaflow-lite/manifests/reduce/uv.lock | 8 + .../pynumaflow_lite/__init__.py | 2 + .../pynumaflow_lite/_reduce_dtypes.py | 6 +- .../pynumaflow_lite/_reduce_server.py | 47 ++++ .../pynumaflow_lite/reducer.pyi | 57 ++++- .../pynumaflow_lite/reducestreamer.pyi | 4 +- packages/pynumaflow-lite/src/reduce/mod.rs | 238 +++++++++++------- packages/pynumaflow-lite/src/reduce/server.rs | 190 +++++++------- packages/pynumaflow-lite/tests/bin/reduce.rs | 4 +- .../tests/examples/reduce_counter_class.py | 55 ++-- .../tests/examples/reduce_counter_func.py | 44 +--- packages/pynumaflow-lite/tests/test_reduce.py | 21 ++ .../pynumaflow-lite/tests/test_reduce_unit.py | 153 +++++++++++ .../pynumaflow-lite/tests/test_sink_unit.py | 7 +- 18 files changed, 616 insertions(+), 361 deletions(-) create mode 100644 packages/pynumaflow-lite/manifests/reduce/.dockerignore create mode 100644 packages/pynumaflow-lite/manifests/reduce/uv.lock create mode 100644 packages/pynumaflow-lite/pynumaflow_lite/_reduce_server.py create mode 100644 packages/pynumaflow-lite/tests/test_reduce_unit.py diff --git a/packages/pynumaflow-lite/manifests/reduce/.dockerignore b/packages/pynumaflow-lite/manifests/reduce/.dockerignore new file mode 100644 index 00000000..21d0b898 --- /dev/null +++ b/packages/pynumaflow-lite/manifests/reduce/.dockerignore @@ -0,0 +1 @@ +.venv/ diff --git a/packages/pynumaflow-lite/manifests/reduce/Dockerfile b/packages/pynumaflow-lite/manifests/reduce/Dockerfile index 1024e5f2..684946fb 100644 --- a/packages/pynumaflow-lite/manifests/reduce/Dockerfile +++ b/packages/pynumaflow-lite/manifests/reduce/Dockerfile @@ -1,39 +1,43 @@ FROM python:3.11-slim-trixie AS builder -ENV PYTHONFAULTHANDLER=1 \ - PYTHONUNBUFFERED=1 \ - PYTHONHASHSEED=random \ - PIP_NO_CACHE_DIR=on \ - PIP_DISABLE_PIP_VERSION_CHECK=on \ - PIP_DEFAULT_TIMEOUT=100 \ - POETRY_HOME="/opt/poetry" \ - POETRY_VIRTUALENVS_IN_PROJECT=true \ - POETRY_NO_INTERACTION=1 \ - PYSETUP_PATH="/opt/pysetup" - -ENV PATH="$POETRY_HOME/bin:$PATH" - -RUN apt-get update \ - && apt-get install --no-install-recommends -y \ - curl \ - wget \ - # deps for building python deps - build-essential \ - && apt-get install -y git \ - && apt-get clean && rm -rf /var/lib/apt/lists/* \ - && curl -sSL https://install.python-poetry.org | python3 - - -FROM builder AS udf - -WORKDIR $PYSETUP_PATH -COPY ./ ./ - -# NOTE: place the built wheel in this directory before building the image -RUN pip install $PYSETUP_PATH/pynumaflow_lite-0.1.0-cp311-cp311-manylinux_2_17_aarch64.manylinux2014_aarch64.whl - -RUN poetry lock -RUN poetry install --no-cache --no-root && \ - rm -rf ~/.cache/pypoetry/ +COPY --from=ghcr.io/astral-sh/uv:0.12.13 /uv /uvx /bin/ -CMD ["python", "reduce_counter_class.py"] +ENV UV_COMPILE_BYTECODE=1 \ + UV_LINK_MODE=copy \ + UV_PYTHON_DOWNLOADS=never + +WORKDIR /app + +RUN --mount=type=cache,target=/root/.cache/uv \ + --mount=type=bind,source=uv.lock,target=uv.lock \ + --mount=type=bind,source=pyproject.toml,target=pyproject.toml \ + uv sync --locked --no-install-project --no-dev + +COPY . . + +RUN --mount=type=cache,target=/root/.cache/uv \ + uv sync --locked --no-dev + +# Install the local pynumaflow-lite wheel that matches the build platform. +RUN uv pip install --no-index --find-links . pynumaflow-lite + +FROM python:3.11-slim-trixie +# Setup a non-root user +RUN groupadd --system --gid 999 nonroot \ + && useradd --system --gid 999 --uid 999 --create-home nonroot + +COPY --from=builder --chown=nonroot:nonroot /app /app + +ENV PATH="/app/.venv/bin:$PATH" + +# Keeps Python from buffering stdout and stderr to avoid situations where +# the application crashes without emitting any logs due to buffering. +ENV PYTHONUNBUFFERED=1 + +# Use the non-root user to run our application +USER nonroot + +WORKDIR /app + +CMD ["python", "reduce_counter_class.py"] diff --git a/packages/pynumaflow-lite/manifests/reduce/pyproject.toml b/packages/pynumaflow-lite/manifests/reduce/pyproject.toml index e36b2026..fee630ae 100644 --- a/packages/pynumaflow-lite/manifests/reduce/pyproject.toml +++ b/packages/pynumaflow-lite/manifests/reduce/pyproject.toml @@ -6,11 +6,6 @@ authors = [ { name = "Vigith Maurice", email = "vigith@gmail.com" } ] readme = "README.md" -requires-python = ">=3.11" +requires-python = "==3.11.*" dependencies = [ ] - -[build-system] -requires = ["poetry-core>=2.0.0,<3.0.0"] -build-backend = "poetry.core.masonry.api" - diff --git a/packages/pynumaflow-lite/manifests/reduce/reduce_counter_class.py b/packages/pynumaflow-lite/manifests/reduce/reduce_counter_class.py index 6f112b4c..169f48f0 100644 --- a/packages/pynumaflow-lite/manifests/reduce/reduce_counter_class.py +++ b/packages/pynumaflow-lite/manifests/reduce/reduce_counter_class.py @@ -1,49 +1,34 @@ import asyncio -import signal +import functools from collections.abc import AsyncIterable -from pynumaflow_lite import reducer +from pynumaflow_lite.reducer import Datum, Message, Metadata, ReduceAsyncServer, Reducer -class ReduceCounter(reducer.Reducer): - def __init__(self, initial: int = 0) -> None: - self.counter = initial - - async def handler( - self, - keys: list[str], - datums: AsyncIterable[reducer.Datum], - md: reducer.Metadata, - ) -> reducer.Messages: - iw = md.interval_window - self.counter = 0 - async for _ in datums: - self.counter += 1 - msg = (f"counter:{self.counter} interval_window_start:{iw.start} interval_window_end:{iw.end}").encode() - out = reducer.Messages() - out.append(reducer.Message(msg, keys)) - return out +class ReduceCounter(Reducer): + """Counts the datums of each window. The server creates a new instance for each window.""" + def __init__(self, initial: int = 0) -> None: + self.initial = initial -async def start(creator: type[reducer.Reducer], init_args: tuple): - sock_file = "/var/run/numaflow/reduce.sock" - server_info_file = "/var/run/numaflow/reducer-server-info" - server = reducer.ReduceAsyncServer(sock_file, server_info_file) + async def handler(self, keys: list[str], datums: AsyncIterable[Datum], md: Metadata) -> list[Message]: + counter = self.initial + async for _ in datums: + counter += 1 + iw = md.interval_window + msg = f"counter:{counter} interval_window_start:{iw.start} interval_window_end:{iw.end}" + return [Message(msg.encode(), keys=keys)] - 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(creator, init_args) - print("Shutting down gracefully...") - except asyncio.CancelledError: - server.stop() - return +async def main() -> None: + # ReduceAsyncServer needs a factory: a callable that takes no arguments and + # returns a new Reducer. The server calls it at the start of every window (once + # for each set of keys), so each window starts with its own state. + # functools.partial binds initial=100, so the server can call + # ReduceCounter(initial=100) without arguments. + # `serve` returns when SIGINT or SIGTERM arrives. + await ReduceAsyncServer(functools.partial(ReduceCounter, initial=100)).serve() if __name__ == "__main__": - asyncio.run(start(ReduceCounter, (0,))) + asyncio.run(main()) diff --git a/packages/pynumaflow-lite/manifests/reduce/uv.lock b/packages/pynumaflow-lite/manifests/reduce/uv.lock new file mode 100644 index 00000000..f16ad9b5 --- /dev/null +++ b/packages/pynumaflow-lite/manifests/reduce/uv.lock @@ -0,0 +1,8 @@ +version = 1 +revision = 3 +requires-python = "==3.11.*" + +[[package]] +name = "reduce-counter" +version = "0.1.0" +source = { virtual = "." } diff --git a/packages/pynumaflow-lite/pynumaflow_lite/__init__.py b/packages/pynumaflow-lite/pynumaflow_lite/__init__.py index 5f8f0c14..2594994b 100644 --- a/packages/pynumaflow-lite/pynumaflow_lite/__init__.py +++ b/packages/pynumaflow-lite/pynumaflow_lite/__init__.py @@ -12,6 +12,7 @@ from ._mapstream_dtypes import MapStreamer from ._mapstream_server import MapStreamAsyncServer from ._reduce_dtypes import Reducer +from ._reduce_server import ReduceAsyncServer from ._reducestreamer_dtypes import ReduceStreamer from ._session_reduce_dtypes import SessionReducer from ._sideinput_dtypes import SideInput @@ -45,6 +46,7 @@ mapstreamer.MapStreamer = MapStreamer mapstreamer.MapStreamAsyncServer = MapStreamAsyncServer reducer.Reducer = Reducer +reducer.ReduceAsyncServer = ReduceAsyncServer session_reducer.SessionReducer = SessionReducer reducestreamer.ReduceStreamer = ReduceStreamer accumulator.Accumulator = Accumulator diff --git a/packages/pynumaflow-lite/pynumaflow_lite/_reduce_dtypes.py b/packages/pynumaflow-lite/pynumaflow_lite/_reduce_dtypes.py index 57aeacb6..947c9fc5 100644 --- a/packages/pynumaflow-lite/pynumaflow_lite/_reduce_dtypes.py +++ b/packages/pynumaflow-lite/pynumaflow_lite/_reduce_dtypes.py @@ -1,7 +1,7 @@ from abc import ABCMeta, abstractmethod from collections.abc import AsyncIterable -from pynumaflow_lite.reducer import Datum, Messages, Metadata +from pynumaflow_lite.reducer import Datum, Message, Metadata class Reducer(metaclass=ABCMeta): @@ -13,8 +13,8 @@ def __call__(self, *args, **kwargs): return self.handler(*args, **kwargs) @abstractmethod - async def handler(self, keys: list[str], datums: AsyncIterable[Datum], md: Metadata) -> Messages: + async def handler(self, keys: list[str], datums: AsyncIterable[Datum], md: Metadata) -> list[Message]: """ - Implement this handler; consume `datums` async iterable and return Messages. + Implement this handler; consume `datums` async iterable and return the messages. """ pass diff --git a/packages/pynumaflow-lite/pynumaflow_lite/_reduce_server.py b/packages/pynumaflow-lite/pynumaflow_lite/_reduce_server.py new file mode 100644 index 00000000..83c21406 --- /dev/null +++ b/packages/pynumaflow-lite/pynumaflow_lite/_reduce_server.py @@ -0,0 +1,47 @@ +from __future__ import annotations + +import inspect +from collections.abc import AsyncIterable, Awaitable, Callable +from typing import TypeAlias + +from ._async_server import _AsyncServer +from ._reduce_dtypes import Reducer +from .pynumaflow_lite import reducer as _reducer + +Datum: TypeAlias = _reducer.Datum +Message: TypeAlias = _reducer.Message +Metadata: TypeAlias = _reducer.Metadata + +ReduceHandler: TypeAlias = Callable[[list[str], AsyncIterable[Datum], Metadata], Awaitable[list[Message]]] + + +class ReduceAsyncServer(_AsyncServer): + """Reduce server. + + `handler` is an async function, or a zero-argument factory that returns a + handler, such as a `Reducer` subclass or `functools.partial(MyReducer, x=1)`. + The server calls a factory once for each window and key set, so each call + gets its own state. + """ + + _name = "reduce" + + def __init__( + self, + handler: ReduceHandler | Callable[[], ReduceHandler], + *, + sock_file: str | None = None, + server_info_file: str | None = None, + install_signal_handlers: bool = True, + ) -> None: + if isinstance(handler, Reducer): + raise TypeError( + "pass the Reducer class or a factory, not an instance: " + "the server creates a new instance for each window" + ) + factory = (lambda: handler) if inspect.iscoroutinefunction(handler) else handler + super().__init__( + _reducer._ReduceAsyncServer(sock_file, server_info_file), + factory, + install_signal_handlers=install_signal_handlers, + ) diff --git a/packages/pynumaflow-lite/pynumaflow_lite/reducer.pyi b/packages/pynumaflow-lite/pynumaflow_lite/reducer.pyi index f12e0f96..87228882 100644 --- a/packages/pynumaflow-lite/pynumaflow_lite/reducer.pyi +++ b/packages/pynumaflow-lite/pynumaflow_lite/reducer.pyi @@ -1,5 +1,7 @@ import datetime as _dt from collections.abc import AsyncIterable, Awaitable, Callable +from types import TracebackType +from typing import Self, TypeAlias # Re-export the Python ABC for user convenience and typing from ._reduce_dtypes import Reducer as Reducer @@ -16,23 +18,26 @@ class Message: tags: list[str] | None = ..., ) -> None: ... @staticmethod - def message_to_drop() -> Message: ... - @staticmethod def to_drop() -> Message: ... - -class Messages: - def __init__(self) -> None: ... - def append(self, message: Message) -> None: ... def __repr__(self) -> str: ... - def __str__(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: ... @@ -40,27 +45,53 @@ class IntervalWindow: start: _dt.datetime end: _dt.datetime + def __init__(self, start: _dt.datetime, end: _dt.datetime) -> None: ... + class Metadata: interval_window: IntervalWindow -class ReduceAsyncServer: + def __init__(self, *, interval_window: IntervalWindow) -> None: ... + +_ReduceHandler: TypeAlias = Callable[[list[str], AsyncIterable[Datum], Metadata], Awaitable[list[Message]]] + +class _ReduceAsyncServer: def __init__( self, sock_file: str | None = ..., - info_file: str | None = ..., + server_info_file: str | None = ..., ) -> None: ... def start( self, - py_creator: (type[Reducer] | Callable[[list[str], AsyncIterable[Datum], Metadata], Awaitable[Messages]]), - init_args: tuple | None = ..., + factory: Callable[[], _ReduceHandler], ) -> Awaitable[None]: ... + def wait_ready(self) -> Awaitable[None]: ... def stop(self) -> None: ... +class ReduceAsyncServer: + def __init__( + self, + handler: _ReduceHandler | Callable[[], _ReduceHandler], + *, + sock_file: str | None = ..., + server_info_file: str | None = ..., + install_signal_handlers: bool = ..., + ) -> None: ... + async def serve(self) -> None: ... + def stop(self) -> None: ... + async def wait_ready(self) -> None: ... + async def wait_for_termination(self) -> None: ... + async def __aenter__(self) -> Self: ... + async def __aexit__( + self, + exc_type: type[BaseException] | None, + exc: BaseException | None, + tb: TracebackType | None, + ) -> None: ... + __all__ = [ "Datum", "IntervalWindow", "Message", - "Messages", "Metadata", "ReduceAsyncServer", "Reducer", diff --git a/packages/pynumaflow-lite/pynumaflow_lite/reducestreamer.pyi b/packages/pynumaflow-lite/pynumaflow_lite/reducestreamer.pyi index 2d28a7f9..92c3244f 100644 --- a/packages/pynumaflow-lite/pynumaflow_lite/reducestreamer.pyi +++ b/packages/pynumaflow-lite/pynumaflow_lite/reducestreamer.pyi @@ -16,13 +16,13 @@ class Message: tags: list[str] | None = ..., ) -> None: ... @staticmethod - def message_to_drop() -> Message: ... + def to_drop() -> Message: ... class Datum: keys: list[str] value: bytes watermark: _dt.datetime - eventtime: _dt.datetime + event_time: _dt.datetime headers: dict[str, str] def __repr__(self) -> str: ... diff --git a/packages/pynumaflow-lite/src/reduce/mod.rs b/packages/pynumaflow-lite/src/reduce/mod.rs index 8bff323e..38715b85 100644 --- a/packages/pynumaflow-lite/src/reduce/mod.rs +++ b/packages/pynumaflow-lite/src/reduce/mod.rs @@ -1,20 +1,35 @@ use chrono::{DateTime, Utc}; +use numaflow::proto::reduce::reduce_client::ReduceClient; use numaflow::reduce; use std::collections::HashMap; +use std::path::PathBuf; use std::sync::Mutex; +use std::time::Duration; pub mod server; use pyo3::prelude::*; +use tokio::net::UnixStream; use tokio::sync::mpsc; +use tonic::transport::Uri; +use tower::service_fn; + +use crate::pyrs::{bytes_literal, py_repr}; /// A message to be sent to the next vertex from a reduce handler. -#[pyclass(module = "pynumaflow_lite.reducer", from_py_object)] -#[derive(Clone, Default, Debug)] +#[pyclass(module = "pynumaflow_lite.reducer", from_py_object, eq)] +#[derive(Clone, Default, Debug, PartialEq)] pub struct Message { - pub keys: Option>, // optional keys - pub value: Vec, // payload - pub tags: Option>, // optional tags (e.g., DROP) + /// 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>, } #[pymethods] @@ -25,16 +40,25 @@ 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![], tags: Some(vec![numaflow::shared::DROP.to_string()]), } } + + fn __repr__(&self, py: Python<'_>) -> PyResult { + Ok(format!( + "Message(value={}, keys={}, tags={})", + bytes_literal(&self.value), + py_repr(py, &self.keys)?, + py_repr(py, &self.tags)?, + )) + } } impl From for reduce::Message { @@ -47,96 +71,79 @@ impl From for reduce::Message { } } -/// A collection of Messages returned by reducer. -#[pyclass(module = "pynumaflow_lite.reducer", 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) - } -} - /// The incoming ReduceRequest accessible in Python function (streamed). #[pyclass(module = "pynumaflow_lite.reducer")] pub struct Datum { + /// Set of keys in the (key, value) terminology of map/reduce paradigm. #[pyo3(get)] pub keys: Vec, + /// The value in the (key, value) terminology of map/reduce paradigm. #[pyo3(get)] pub value: Vec, + /// [watermark](https://numaflow.numaproj.io/core-concepts/watermarks/) represented by time is a + /// guarantee that we will not see an element older than this time. #[pyo3(get)] 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, } +#[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: Vec, - value: Vec, - watermark: DateTime, - eventtime: DateTime, - headers: HashMap, + keys: Option>, + value: Option>, + event_time: Option>, + watermark: Option>, + headers: Option>, ) -> Self { Self { - keys, - value, - watermark, - eventtime, - 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), + headers: headers.unwrap_or_default(), } } - fn __repr__(&self) -> String { - format!( - "Datum(keys={:?}, value={:?}, watermark={}, eventtime={}, headers={:?})", - self.keys, self.value, self.watermark, self.eventtime, self.headers - ) + fn __repr__(&self, py: Python<'_>) -> PyResult { + Ok(format!( + "Datum(keys={}, value={}, watermark={}, event_time={}, headers={})", + py_repr(py, &self.keys)?, + bytes_literal(&self.value), + self.watermark, + self.event_time, + py_repr(py, &self.headers)?, + )) } - fn __str__(&self) -> String { - format!( - "Datum(keys={:?}, value={:?}, watermark={}, eventtime={}, headers={:?})", - self.keys, - String::from_utf8_lossy(&self.value), - self.watermark, - self.eventtime, - self.headers - ) + fn __str__(&self, py: Python<'_>) -> PyResult { + self.__repr__(py) } } impl From for Datum { fn from(value: reduce::ReduceRequest) -> Self { - Self::new( - value.keys, - value.value, - value.watermark, - value.eventtime, - value.headers, - ) + Self { + keys: value.keys, + value: value.value, + watermark: value.watermark, + event_time: value.eventtime, + headers: value.headers, + } } } @@ -150,7 +157,10 @@ pub struct IntervalWindow { pub end: DateTime, } +#[pymethods] impl IntervalWindow { + #[new] + #[pyo3(signature = (start: "datetime.datetime", end: "datetime.datetime") -> "IntervalWindow")] pub(crate) fn new(start: DateTime, end: DateTime) -> Self { Self { start, end } } @@ -164,7 +174,10 @@ pub struct Metadata { pub interval_window: IntervalWindow, } +#[pymethods] impl Metadata { + #[new] + #[pyo3(signature = (*, interval_window: "IntervalWindow") -> "Metadata")] pub(crate) fn new(interval_window: IntervalWindow) -> Self { Self { interval_window } } @@ -204,48 +217,89 @@ impl PyAsyncDatumStream { } } -/// Async Reduce Server that can be started from Python code, taking a class (creator). -#[pyclass(module = "pynumaflow_lite.reducer")] +async fn reduce_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(ReduceClient::new(channel)) +} + +// Poll the IsReady probe until it succeeds. There is no deadline: Python sets +// one with asyncio.timeout(), and cancellation drops this future. +async fn wait_for_ready(sock_file: String) { + loop { + if let Ok(mut client) = reduce_client(sock_file.clone()).await + && let Ok(response) = client.is_ready(()).await + && response.into_inner().ready + { + return; + } + + tokio::time::sleep(Duration::from_millis(100)).await; + } +} + +/// Async Reduce Server that can be started from Python code, taking a reducer factory. +#[pyclass(name = "_ReduceAsyncServer", module = "pynumaflow_lite.reducer")] pub struct ReduceAsyncServer { sock_file: String, - info_file: String, + server_info_file: String, shutdown_tx: Mutex>>, } #[pymethods] impl ReduceAsyncServer { #[new] - #[pyo3(signature = (sock_file: "str | None"=reduce::SOCK_ADDR.to_string(), info_file: "str | None"=reduce::SERVER_INFO_FILE.to_string()) -> "ReduceAsyncServer")] - fn new(sock_file: String, info_file: String) -> Self { + #[pyo3(signature = ( + sock_file: "str | None"=None, + server_info_file: "str | None"=None, + ) -> "_ReduceAsyncServer")] + fn new(sock_file: Option, server_info_file: Option) -> Self { Self { - sock_file, - info_file, + sock_file: sock_file.unwrap_or_else(|| reduce::SOCK_ADDR.to_string()), + server_info_file: server_info_file + .unwrap_or_else(|| reduce::SERVER_INFO_FILE.to_string()), shutdown_tx: Mutex::new(None), } } - /// Start the server with the given Python class (creator) or function. - /// - For class-based: pass the class and optionally init_args (tuple). - /// - For function-based: pass the async function directly (init_args is ignored). - #[pyo3(signature = (py_creator: "object", init_args: "object | None"=None) -> "None")] - pub fn start<'a>( - &self, - py: Python<'a>, - py_creator: Py, - init_args: Option>, - ) -> PyResult> { + /// Start the server with the given zero-argument factory. The server calls the + /// factory once for each window and key set to get the async handler. + #[pyo3(signature = (factory: "callable") -> "None")] + pub fn start<'a>(&self, py: Python<'a>, factory: 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(); *guard = Some(tx); } + pyo3_async_runtimes::tokio::future_into_py( + py, + crate::reduce::server::start(factory, sock_file, server_info_file, rx), + ) + } + + /// Wait until the Numaflow IsReady probe succeeds over the reduce UDS. + #[pyo3(signature = () -> "None")] + pub fn wait_ready<'a>(&self, py: Python<'a>) -> PyResult> { + let sock_file = self.sock_file.clone(); + pyo3_async_runtimes::tokio::future_into_py(py, async move { - crate::reduce::server::start(py_creator, init_args, sock_file, info_file, rx) - .await - .expect("server failed to start"); + wait_for_ready(sock_file).await; Ok(()) }) } @@ -263,11 +317,9 @@ impl ReduceAsyncServer { /// Helper to populate a PyModule with reduce 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::()?; - m.add_class::()?; Ok(()) } diff --git a/packages/pynumaflow-lite/src/reduce/server.rs b/packages/pynumaflow-lite/src/reduce/server.rs index 68d35762..b96834d4 100644 --- a/packages/pynumaflow-lite/src/reduce/server.rs +++ b/packages/pynumaflow-lite/src/reduce/server.rs @@ -1,64 +1,60 @@ +use std::sync::{Arc, Mutex}; + use crate::reduce::{ Datum as PyDatum, IntervalWindow as PyIntervalWindow, Message as PyMessage, - Messages as PyMessages, Metadata as PyMetadata, PyAsyncDatumStream, + Metadata as PyMetadata, PyAsyncDatumStream, }; use numaflow::reduce; use numaflow::shared::ServerExtras; +use pyo3::exceptions::PyTypeError; use pyo3::prelude::*; -use pyo3::types::PyTuple; -use std::sync::Arc; + +use crate::pyrs::{combine_errors, format_error}; pub(crate) struct PyReduceCreator { /// handle to Python event loop pub(crate) event_loop: Arc>, - /// Python class to instantiate per window OR a function - pub(crate) py_creator: Arc>, - /// optional tuple of positional args (only for class) - pub(crate) init_args: Option>>, - /// true if py_creator is a function, false if it's a class. - pub(crate) is_function: bool, + /// Zero-argument Python callable that returns the handler for one window + pub(crate) factory: Arc>, + pub(crate) errors: Arc>>, } pub(crate) struct PyReduceRunner { /// handle to Python event loop pub(crate) event_loop: Arc>, - /// Instance of class per window OR the function itself. It can be class because we have overloaded - /// __call__ method to make it callable like a function. - pub(crate) py_obj: Arc>, + pub(crate) factory: Arc>, + pub(crate) errors: Arc>>, } -#[tonic::async_trait] impl reduce::ReducerCreator for PyReduceCreator { type R = PyReduceRunner; + // The factory runs later, in reduce(). A panic in create() would stop the numaflow + // task executor without an error to the client, but reduce() reports a panic. fn create(&self) -> Self::R { - // If it's a function, just clone the reference; if it's a class, instantiate it. - let inst = if self.is_function { - // For functions, we don't instantiate; just use the function directly - Python::attach(|py| self.py_creator.clone_ref(py)) - } else { - // For classes, instantiate the Python class synchronously under the GIL. - Python::attach(|py| { - let class = self.py_creator.as_ref(); - match &self.init_args { - Some(args) => { - let bound = args.as_ref().bind(py); - let py_tuple = bound.cast::()?; - class.call1(py, py_tuple) - } - None => class.call0(py), - } - }) - .expect("failed to instantiate Python reducer class") - }; - PyReduceRunner { event_loop: self.event_loop.clone(), - py_obj: Arc::new(inst), + factory: self.factory.clone(), + errors: self.errors.clone(), } } } +impl PyReduceRunner { + fn fail(&self, error: PyErr) -> ! { + // numaflow calls reduce() concurrently, so each error belongs to a + // different window. Keep all of them. start() raises them together, which + // lets Python format every traceback instead of Rust printing them by hand. + let message = Python::attach(|py| format_error(py, &error)); + self.errors.lock().unwrap().push(error); + + // numaflow catches this panic, sends a gRPC error for this window, and + // starts the server shutdown. An empty result would instead look like a + // window that the handler reduced to no messages on purpose. + panic!("{message}"); + } +} + #[tonic::async_trait] impl reduce::Reducer for PyReduceRunner { async fn reduce( @@ -73,97 +69,96 @@ impl reduce::Reducer for PyReduceRunner { // Spawn a task forwarding incoming datums to the Python-facing channel let forwarder = tokio::spawn(async move { while let Some(req) = input.recv().await { - let datum = PyDatum::from(req); - if tx.send(datum).await.is_err() { + if tx.send(req.into()).await.is_err() { break; } } }); // Call the Python coroutine: - // - For class-based: obj.handler(keys, datums: AsyncIterable[Datum], md: Metadata) -> Messages - // - For function-based: obj(keys, datums: AsyncIterable[Datum], md: Metadata) -> Messages - let fut = Python::attach(|py| { + // factory()(keys: list[str], datums: AsyncIterable[Datum], md: Metadata) -> list[Message] + let fut = match Python::attach(|py| -> PyResult<_> { let locals = pyo3_async_runtimes::TaskLocals::new(self.event_loop.bind(py).clone()); - let obj = self.py_obj.clone(); + let handler = self.factory.call0(py)?; - // Build metadata and stream objects let stream = PyAsyncDatumStream::new_with(rx); let interval = PyIntervalWindow::new(md.interval_window.start_time, md.interval_window.end_time); let md_py = PyMetadata::new(interval); - // doesn't matter whether it is func or class, obj is callable - let coro = obj - .call1(py, (keys.clone(), stream, md_py)) - .unwrap() - .into_bound(py); - - pyo3_async_runtimes::into_future_with_locals(&locals, coro).unwrap() - }); + let coro = handler.call1(py, (keys, stream, md_py))?.into_bound(py); + let is_awaitable: bool = py + .import("inspect")? + .call_method1("isawaitable", (&coro,))? + .extract()?; + if !is_awaitable { + return Err(PyErr::new::( + "reduce handler must be an async function (coroutine)", + )); + } + pyo3_async_runtimes::into_future_with_locals(&locals, coro).map_err(|_| { + PyErr::new::("reduce handler must be an async function (coroutine)") + }) + }) { + Ok(fut) => fut, + Err(error) => self.fail(error), + }; - let result = fut.await.unwrap(); + let result = match fut.await { + Ok(result) => result, + Err(error) => self.fail(error), + }; // Ensure forwarder completes let _ = forwarder.await; - Python::attach(|py| { - // Expect Messages; also allow a single Message for convenience - if let Ok(msgs) = result.extract::(py) { - msgs.messages - .into_iter() - .map(reduce::Message::from) - .collect::>() - } else if let Ok(single) = result.extract::(py) { - vec![single.into()] - } else { - vec![] - } - }) + 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!( + "reduce handler must return list[Message], got {type_name}" + )) + }) + }) { + Ok(messages) => messages, + Err(error) => self.fail(error), + }; + + messages.into_iter().map(|m| m.into()).collect() } } /// Start the reduce server by spinning up a dedicated Python asyncio loop and wiring shutdown. pub(super) async fn start( - py_creator: Py, - init_args: Option>, + factory: Py, sock_file: String, info_file: String, shutdown_rx: tokio::sync::oneshot::Receiver<()>, ) -> Result<(), pyo3::PyErr> { let (tx, rx) = tokio::sync::oneshot::channel(); - let py_asyncio_loop_handle = tokio::task::spawn_blocking(move || crate::pyrs::run_asyncio(tx)); + let py_asyncio_loop_handle = tokio::task::spawn_blocking({ + println!( + "Starting Reduce 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); - - // Detect if py_creator is a function or a class - let is_function = Python::attach(|py| { - let obj = py_creator.bind(py); - // Check if it's a function or coroutine function using inspect module - let inspect = py.import("inspect").ok()?; - if let Ok(is_func) = inspect.call_method1("isfunction", (obj,)) - && let Ok(result) = is_func.extract::() - && result - { - return Some(true); - } - // Also check for coroutine function - if let Ok(is_coro) = inspect.call_method1("iscoroutinefunction", (obj,)) - && let Ok(result) = is_coro.extract::() - && result - { - return Some(true); - } - Some(false) - }) - .unwrap_or(false); + let errors = Arc::new(Mutex::new(Vec::new())); + // Shutdown has two sources, and neither one needs a channel here. The Python + // side signals stop() through shutdown_rx. An uncaught Python error panics in + // fail(), and numaflow then shuts the server down on its own. let creator = PyReduceCreator { event_loop: event_loop.clone(), - py_creator: Arc::new(py_creator), - init_args: init_args.map(Arc::new), - is_function, + factory: Arc::new(factory), + errors: Arc::clone(&errors), }; let server = reduce::Server::new(creator) @@ -171,7 +166,7 @@ pub(super) async fn start( .with_server_info_file(info_file); let result = server - .start_with_shutdown(combined_rx) + .start_with_shutdown(shutdown_rx) .await .map_err(|e| pyo3::PyErr::new::(e.to_string())); @@ -182,15 +177,14 @@ pub(super) async fn start( } }); - println!("Numaflow Core (reduce) has shutdown..."); + println!("Numaflow Reduce has shutdown..."); // Wait for the blocking asyncio thread to finish. let _ = py_asyncio_loop_handle.await; - // if not finished, abort it - if !sig_handle.is_finished() { - println!("Aborting signal handler"); - sig_handle.abort(); + let errors = std::mem::take(&mut *errors.lock().unwrap()); + if !errors.is_empty() { + return Err(Python::attach(|py| combine_errors(py, errors))); } result diff --git a/packages/pynumaflow-lite/tests/bin/reduce.rs b/packages/pynumaflow-lite/tests/bin/reduce.rs index 9174da8a..353fb5c1 100644 --- a/packages/pynumaflow-lite/tests/bin/reduce.rs +++ b/packages/pynumaflow-lite/tests/bin/reduce.rs @@ -75,7 +75,9 @@ async fn main() -> Result<(), Box> { loop { if let Some(r) = resp.message().await? { if let Some(res) = r.result { - assert!(!res.value.is_empty()); + let value = String::from_utf8(res.value)?; + assert!(value.starts_with("counter:1 "), "unexpected value: {value}"); + assert_eq!(res.keys, vec!["k".to_string()]); break; } } else { diff --git a/packages/pynumaflow-lite/tests/examples/reduce_counter_class.py b/packages/pynumaflow-lite/tests/examples/reduce_counter_class.py index c9783973..92b632be 100644 --- a/packages/pynumaflow-lite/tests/examples/reduce_counter_class.py +++ b/packages/pynumaflow-lite/tests/examples/reduce_counter_class.py @@ -1,49 +1,30 @@ import asyncio -import signal +import functools from collections.abc import AsyncIterable -from pynumaflow_lite import reducer +from pynumaflow_lite.reducer import Datum, Message, Metadata, ReduceAsyncServer, Reducer -class ReduceCounter(reducer.Reducer): +class ReduceCounter(Reducer): def __init__(self, initial: int = 0) -> None: - self.counter = initial - - async def handler( - self, - keys: list[str], - datums: AsyncIterable[reducer.Datum], - md: reducer.Metadata, - ) -> reducer.Messages: - iw = md.interval_window - self.counter = 0 - async for _ in datums: - self.counter += 1 - msg = (f"counter:{self.counter} interval_window_start:{iw.start} interval_window_end:{iw.end}").encode() - out = reducer.Messages() - out.append(reducer.Message(msg, keys)) - return out - + self.initial = initial -async def start(creator: type, init_args: tuple): - sock_file = "/tmp/var/run/numaflow/reduce.sock" - server_info_file = "/tmp/var/run/numaflow/reducer-server-info" - server = reducer.ReduceAsyncServer(sock_file, server_info_file) + async def handler(self, keys: list[str], datums: AsyncIterable[Datum], md: Metadata) -> list[Message]: + counter = self.initial + async for _ in datums: + counter += 1 + iw = md.interval_window + msg = f"counter:{counter} interval_window_start:{iw.start} interval_window_end:{iw.end}" + return [Message(msg.encode(), keys=keys)] - 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(creator, init_args) - print("Shutting down gracefully...") - except asyncio.CancelledError: - server.stop() - return +async def main() -> None: + await ReduceAsyncServer( + functools.partial(ReduceCounter, initial=0), + sock_file="/tmp/var/run/numaflow/reduce.sock", + server_info_file="/tmp/var/run/numaflow/reducer-server-info", + ).serve() if __name__ == "__main__": - asyncio.run(start(ReduceCounter, (0,))) + asyncio.run(main()) diff --git a/packages/pynumaflow-lite/tests/examples/reduce_counter_func.py b/packages/pynumaflow-lite/tests/examples/reduce_counter_func.py index 8f67680c..af168d6b 100644 --- a/packages/pynumaflow-lite/tests/examples/reduce_counter_func.py +++ b/packages/pynumaflow-lite/tests/examples/reduce_counter_func.py @@ -1,47 +1,25 @@ import asyncio -import signal -from collections.abc import AsyncIterable, Awaitable, Callable +from collections.abc import AsyncIterable -from pynumaflow_lite import reducer +from pynumaflow_lite.reducer import Datum, Message, Metadata, ReduceAsyncServer -async def reduce_handler( - keys: list[str], datums: AsyncIterable[reducer.Datum], md: reducer.Metadata -) -> reducer.Messages: +async def reduce_handler(keys: list[str], datums: AsyncIterable[Datum], md: Metadata) -> list[Message]: interval_window = md.interval_window counter = 0 async for _ in datums: counter += 1 msg = f"counter:{counter} interval_window_start:{interval_window.start} interval_window_end:{interval_window.end}" - out = reducer.Messages() - out.append(reducer.Message(str.encode(msg), keys=keys)) - return out + return [Message(msg.encode(), keys=keys)] -async def start( - handler: Callable[ - [list[str], AsyncIterable[reducer.Datum], reducer.Metadata], - Awaitable[reducer.Messages], - ], -): - sock_file = "/tmp/var/run/numaflow/reduce.sock" - server_info_file = "/tmp/var/run/numaflow/reducer-server-info" - server = reducer.ReduceAsyncServer(sock_file, server_info_file) - - 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(handler) - print("Shutting down gracefully...") - except asyncio.CancelledError: - server.stop() - return +async def main() -> None: + await ReduceAsyncServer( + reduce_handler, + sock_file="/tmp/var/run/numaflow/reduce.sock", + server_info_file="/tmp/var/run/numaflow/reducer-server-info", + ).serve() if __name__ == "__main__": - asyncio.run(start(reduce_handler)) + asyncio.run(main()) diff --git a/packages/pynumaflow-lite/tests/test_reduce.py b/packages/pynumaflow-lite/tests/test_reduce.py index 8bc2ecca..d72c2be3 100644 --- a/packages/pynumaflow-lite/tests/test_reduce.py +++ b/packages/pynumaflow-lite/tests/test_reduce.py @@ -3,6 +3,7 @@ import pytest from _test_utils import run_python_server_with_rust_client +from pynumaflow_lite import reducer SOCK_PATH = Path("/tmp/var/run/numaflow/reduce.sock") SERVER_INFO = Path("/tmp/var/run/numaflow/reducer-server-info") @@ -21,3 +22,23 @@ def test_python_reduce_server_and_rust_client(script: str, tmp_path: Path): server_info_path=SERVER_INFO, rust_bin_name="test_reduce", ) + + +def test_reduce_data_types_match_pythonic_api(): + datum = reducer.Datum(keys=["key"], value=b"value") + assert datum.keys == ["key"] + assert datum.value == b"value" + assert datum.event_time + assert "value=b'value'" in repr(datum) + + message = reducer.Message(b"value", keys=["key"], tags=["tag"]) + assert message.value == b"value" + assert message.keys == ["key"] + assert message.tags == ["tag"] + assert message == reducer.Message(b"value", keys=["key"], tags=["tag"]) + assert repr(message) == "Message(value=b'value', keys=['key'], tags=['tag'])" + assert reducer.Message.to_drop().tags == ["U+005C__DROP__"] + + assert not hasattr(reducer, "Messages") + assert not hasattr(reducer.Message, "message_to_drop") + assert not hasattr(reducer, "PyAsyncDatumStream") diff --git a/packages/pynumaflow-lite/tests/test_reduce_unit.py b/packages/pynumaflow-lite/tests/test_reduce_unit.py new file mode 100644 index 00000000..b0a6b45f --- /dev/null +++ b/packages/pynumaflow-lite/tests/test_reduce_unit.py @@ -0,0 +1,153 @@ +import asyncio +import functools +import os +import subprocess +import uuid +from collections.abc import AsyncIterable, AsyncIterator +from contextlib import suppress +from datetime import UTC, datetime +from pathlib import Path + +import pytest + +from pynumaflow_lite import reducer + +CARGO_ROOT = Path(__file__).resolve().parent.parent + +WINDOW = reducer.IntervalWindow(datetime(2023, 6, 1, tzinfo=UTC), datetime(2023, 6, 1, 0, 1, tzinfo=UTC)) + + +class Counter(reducer.Reducer): + def __init__(self, initial: int = 0) -> None: + self.counter = initial + + async def handler( + self, keys: list[str], datums: AsyncIterable[reducer.Datum], md: reducer.Metadata + ) -> list[reducer.Message]: + async for _ in datums: + self.counter += 1 + return [reducer.Message(f"counter:{self.counter} ".encode(), keys=keys)] + + +async def _datums(count: int) -> AsyncIterator[reducer.Datum]: + for _ in range(count): + yield reducer.Datum(keys=["key"], value=b"value") + + +def test_reduce_handler_can_be_unit_tested_without_server(): + md = reducer.Metadata(interval_window=WINDOW) + + messages = asyncio.run(Counter(initial=10).handler(["key"], _datums(2), md)) + + assert messages == [reducer.Message(b"counter:12 ", keys=["key"])] + assert md.interval_window.start == WINDOW.start + + +def test_reduce_server_rejects_reducer_instance(): + with pytest.raises(TypeError, match="not an instance"): + reducer.ReduceAsyncServer(Counter()) + + +async def _run_reduce_client(sock_path: Path) -> subprocess.CompletedProcess[str]: + return await asyncio.to_thread( + subprocess.run, + ["cargo", "run", "--quiet", "--bin", "test_reduce", "--", str(sock_path)], + cwd=CARGO_ROOT, + capture_output=True, + text=True, + timeout=60, + check=False, + ) + + +async def _exercise_server( + tmp_path: Path, + handler, + clients: list[subprocess.CompletedProcess[str]] | None = None, +) -> 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 = reducer.ReduceAsyncServer( + handler, + sock_file=str(sock_path), + server_info_file=str(server_info_path), + ) + + try: + async with server: + client = await _run_reduce_client(sock_path) + if clients is not None: + clients.append(client) + finally: + for path in (sock_path, server_info_path): + with suppress(FileNotFoundError): + os.unlink(path) + + +def test_reduce_server_creates_reducer_from_factory(tmp_path: Path): + # The test client sends one datum and checks for "counter:1 ". + clients: list[subprocess.CompletedProcess[str]] = [] + asyncio.run(_exercise_server(tmp_path, functools.partial(Counter, initial=0), clients)) + + [client] = clients + assert client.returncode == 0, client.stderr + + +def test_reduce_server_propagates_handler_exception(tmp_path: Path): + async def handler( + keys: list[str], datums: AsyncIterable[reducer.Datum], md: reducer.Metadata + ) -> list[reducer.Message]: + async for _datum in datums: + raise RuntimeError("reduce exploded") + return [] + + clients: list[subprocess.CompletedProcess[str]] = [] + with pytest.raises(RuntimeError, match="reduce exploded"): + asyncio.run(_exercise_server(tmp_path, handler, clients)) + + # The client must get a gRPC error, not an empty result. Do not check for the + # handler text: numaflow-rs keeps only the first panic of the process. + [client] = clients + assert client.returncode != 0 + assert "UDF_EXECUTION_ERROR" in client.stderr + + +def test_reduce_server_propagates_factory_exception(tmp_path: Path): + def factory(): + raise RuntimeError("factory exploded") + + with pytest.raises(RuntimeError, match="factory exploded"): + asyncio.run(_exercise_server(tmp_path, factory)) + + +def test_reduce_server_rejects_non_list_response(tmp_path: Path): + async def handler(keys: list[str], datums: AsyncIterable[reducer.Datum], md: reducer.Metadata) -> str: + async for _datum in datums: + pass + return "not messages" + + with pytest.raises(TypeError, match=r"reduce handler must return list\[Message\]"): + asyncio.run(_exercise_server(tmp_path, handler)) + + +def test_reduce_server_rejects_sync_handler(tmp_path: Path): + def factory(): + def handler(keys, datums, md): + return [] + + return handler + + with pytest.raises(TypeError, match="reduce handler must be an async function"): + asyncio.run(_exercise_server(tmp_path, factory)) + + +def test_reduce_server_wait_ready_obeys_asyncio_timeout(tmp_path: Path): + server = reducer.ReduceAsyncServer(Counter, sock_file=str(tmp_path / "nobody-listens.sock")) + + async def wait() -> None: + async with asyncio.timeout(0.2): + await server.wait_ready() + + with pytest.raises(TimeoutError): + asyncio.run(wait()) diff --git a/packages/pynumaflow-lite/tests/test_sink_unit.py b/packages/pynumaflow-lite/tests/test_sink_unit.py index 9fea7b25..d369a44b 100644 --- a/packages/pynumaflow-lite/tests/test_sink_unit.py +++ b/packages/pynumaflow-lite/tests/test_sink_unit.py @@ -84,11 +84,12 @@ async def handler(datums: AsyncIterator[sinker.Datum]) -> list[sinker.Response]: with pytest.raises(RuntimeError, match="sink exploded"): asyncio.run(_exercise_server(tmp_path, handler, clients)) - # The client must get a gRPC error that carries the handler error, not an - # empty batch result. + # The client must get a gRPC error, not an empty batch result. Do not check for + # the handler text: numaflow-rs keeps only the first panic of the process, and + # an earlier test in the same pytest process can own it. [client] = clients assert client.returncode != 0 - assert "sink exploded" in client.stderr + assert "UDF_EXECUTION_ERROR" in client.stderr assert "results.len()" not in client.stderr From 421e5f6a07578ca8a070857be1b84bc220efe95c Mon Sep 17 00:00:00 2001 From: Sreekanth Date: Wed, 14 Oct 2026 07:49:26 +0530 Subject: [PATCH 2/2] pynumaflow-lite: Pythonic APIs for ReduceStreamer Signed-off-by: Sreekanth --- .../manifests/batchmap/batchmap_cat.py | 4 +- .../manifests/reduce/reduce_counter_class.py | 4 +- .../manifests/reducestream/.dockerignore | 1 + .../manifests/reducestream/Dockerfile | 74 ++++---- .../manifests/reducestream/pyproject.toml | 7 +- .../reducestream/reducestream_counter.py | 62 +++---- .../manifests/reducestream/uv.lock | 8 + .../pynumaflow_lite/__init__.py | 2 + .../pynumaflow_lite/_batchmap_server.py | 4 +- .../pynumaflow_lite/_batchmapper_dtypes.py | 4 +- .../pynumaflow_lite/_reduce_dtypes.py | 6 +- .../pynumaflow_lite/_reduce_server.py | 4 +- .../pynumaflow_lite/_reducestream_server.py | 47 ++++++ .../pynumaflow_lite/_reducestreamer_dtypes.py | 9 +- .../pynumaflow_lite/batchmapper.pyi | 6 +- .../pynumaflow_lite/reducer.pyi | 4 +- .../pynumaflow_lite/reducestreamer.pyi | 60 +++++-- .../pynumaflow-lite/src/batchmap/server.rs | 2 +- packages/pynumaflow-lite/src/reduce/mod.rs | 13 +- packages/pynumaflow-lite/src/reduce/server.rs | 2 +- .../pynumaflow-lite/src/reducestream/mod.rs | 56 +++--- .../src/reducestream/server.rs | 158 +++++++++-------- .../tests/examples/batchmap_cat.py | 4 +- .../tests/examples/batchmap_cat_class.py | 4 +- .../tests/examples/reduce_counter_class.py | 4 +- .../tests/examples/reduce_counter_func.py | 4 +- .../tests/examples/reducestream_counter.py | 48 ++---- .../examples/reducestream_counter_func.py | 25 +++ .../pynumaflow-lite/tests/test_reduce_unit.py | 8 +- .../tests/test_reducestream.py | 15 ++ .../tests/test_reducestream_unit.py | 159 ++++++++++++++++++ 31 files changed, 542 insertions(+), 266 deletions(-) create mode 100644 packages/pynumaflow-lite/manifests/reducestream/.dockerignore create mode 100644 packages/pynumaflow-lite/manifests/reducestream/uv.lock create mode 100644 packages/pynumaflow-lite/pynumaflow_lite/_reducestream_server.py create mode 100644 packages/pynumaflow-lite/tests/examples/reducestream_counter_func.py create mode 100644 packages/pynumaflow-lite/tests/test_reducestream_unit.py diff --git a/packages/pynumaflow-lite/manifests/batchmap/batchmap_cat.py b/packages/pynumaflow-lite/manifests/batchmap/batchmap_cat.py index e6fd73b6..81f063f7 100644 --- a/packages/pynumaflow-lite/manifests/batchmap/batchmap_cat.py +++ b/packages/pynumaflow-lite/manifests/batchmap/batchmap_cat.py @@ -1,11 +1,11 @@ import asyncio -from collections.abc import AsyncIterable +from collections.abc import AsyncIterator from pynumaflow_lite.batchmapper import BatchMapAsyncServer, BatchMapper, BatchResponse, Datum, Message class SimpleBatchCat(BatchMapper): - async def handler(self, batch: AsyncIterable[Datum]) -> list[BatchResponse]: + async def handler(self, batch: AsyncIterator[Datum]) -> list[BatchResponse]: responses = [] async for datum in batch: if datum.value == b"bad world": diff --git a/packages/pynumaflow-lite/manifests/reduce/reduce_counter_class.py b/packages/pynumaflow-lite/manifests/reduce/reduce_counter_class.py index 169f48f0..1f7cd460 100644 --- a/packages/pynumaflow-lite/manifests/reduce/reduce_counter_class.py +++ b/packages/pynumaflow-lite/manifests/reduce/reduce_counter_class.py @@ -1,6 +1,6 @@ import asyncio import functools -from collections.abc import AsyncIterable +from collections.abc import AsyncIterator from pynumaflow_lite.reducer import Datum, Message, Metadata, ReduceAsyncServer, Reducer @@ -11,7 +11,7 @@ class ReduceCounter(Reducer): def __init__(self, initial: int = 0) -> None: self.initial = initial - async def handler(self, keys: list[str], datums: AsyncIterable[Datum], md: Metadata) -> list[Message]: + async def handler(self, keys: list[str], datums: AsyncIterator[Datum], md: Metadata) -> list[Message]: counter = self.initial async for _ in datums: counter += 1 diff --git a/packages/pynumaflow-lite/manifests/reducestream/.dockerignore b/packages/pynumaflow-lite/manifests/reducestream/.dockerignore new file mode 100644 index 00000000..21d0b898 --- /dev/null +++ b/packages/pynumaflow-lite/manifests/reducestream/.dockerignore @@ -0,0 +1 @@ +.venv/ diff --git a/packages/pynumaflow-lite/manifests/reducestream/Dockerfile b/packages/pynumaflow-lite/manifests/reducestream/Dockerfile index 60cd2234..aeb06d36 100644 --- a/packages/pynumaflow-lite/manifests/reducestream/Dockerfile +++ b/packages/pynumaflow-lite/manifests/reducestream/Dockerfile @@ -1,39 +1,43 @@ FROM python:3.11-slim-trixie AS builder -ENV PYTHONFAULTHANDLER=1 \ - PYTHONUNBUFFERED=1 \ - PYTHONHASHSEED=random \ - PIP_NO_CACHE_DIR=on \ - PIP_DISABLE_PIP_VERSION_CHECK=on \ - PIP_DEFAULT_TIMEOUT=100 \ - POETRY_HOME="/opt/poetry" \ - POETRY_VIRTUALENVS_IN_PROJECT=true \ - POETRY_NO_INTERACTION=1 \ - PYSETUP_PATH="/opt/pysetup" - -ENV PATH="$POETRY_HOME/bin:$PATH" - -RUN apt-get update \ - && apt-get install --no-install-recommends -y \ - curl \ - wget \ - # deps for building python deps - build-essential \ - && apt-get install -y git \ - && apt-get clean && rm -rf /var/lib/apt/lists/* \ - && curl -sSL https://install.python-poetry.org | python3 - - -FROM builder AS udf - -WORKDIR $PYSETUP_PATH -COPY ./ ./ - -# NOTE: place the built wheel in this directory before building the image -RUN pip install $PYSETUP_PATH/pynumaflow_lite-0.1.0-cp311-cp311-manylinux_2_17_aarch64.manylinux2014_aarch64.whl - -RUN poetry lock -RUN poetry install --no-cache --no-root && \ - rm -rf ~/.cache/pypoetry/ +COPY --from=ghcr.io/astral-sh/uv:0.12.13 /uv /uvx /bin/ -CMD ["python", "reducestream_counter.py"] +ENV UV_COMPILE_BYTECODE=1 \ + UV_LINK_MODE=copy \ + UV_PYTHON_DOWNLOADS=never + +WORKDIR /app + +RUN --mount=type=cache,target=/root/.cache/uv \ + --mount=type=bind,source=uv.lock,target=uv.lock \ + --mount=type=bind,source=pyproject.toml,target=pyproject.toml \ + uv sync --locked --no-install-project --no-dev + +COPY . . + +RUN --mount=type=cache,target=/root/.cache/uv \ + uv sync --locked --no-dev + +# Install the local pynumaflow-lite wheel that matches the build platform. +RUN uv pip install --no-index --find-links . pynumaflow-lite + +FROM python:3.11-slim-trixie +# Setup a non-root user +RUN groupadd --system --gid 999 nonroot \ + && useradd --system --gid 999 --uid 999 --create-home nonroot + +COPY --from=builder --chown=nonroot:nonroot /app /app + +ENV PATH="/app/.venv/bin:$PATH" + +# Keeps Python from buffering stdout and stderr to avoid situations where +# the application crashes without emitting any logs due to buffering. +ENV PYTHONUNBUFFERED=1 + +# Use the non-root user to run our application +USER nonroot + +WORKDIR /app + +CMD ["python", "reducestream_counter.py"] diff --git a/packages/pynumaflow-lite/manifests/reducestream/pyproject.toml b/packages/pynumaflow-lite/manifests/reducestream/pyproject.toml index c5eb1a17..50bb8d7e 100644 --- a/packages/pynumaflow-lite/manifests/reducestream/pyproject.toml +++ b/packages/pynumaflow-lite/manifests/reducestream/pyproject.toml @@ -6,11 +6,6 @@ authors = [ { name = "Vigith Maurice", email = "vigith@gmail.com" } ] readme = "README.md" -requires-python = ">=3.11" +requires-python = "==3.11.*" dependencies = [ ] - -[build-system] -requires = ["poetry-core>=2.0.0,<3.0.0"] -build-backend = "poetry.core.masonry.api" - diff --git a/packages/pynumaflow-lite/manifests/reducestream/reducestream_counter.py b/packages/pynumaflow-lite/manifests/reducestream/reducestream_counter.py index 0097df6c..87368845 100644 --- a/packages/pynumaflow-lite/manifests/reducestream/reducestream_counter.py +++ b/packages/pynumaflow-lite/manifests/reducestream/reducestream_counter.py @@ -9,36 +9,31 @@ """ import asyncio -import signal -from collections.abc import AsyncIterable, AsyncIterator +import functools +from collections.abc import AsyncIterator -from pynumaflow_lite import reducestreamer +from pynumaflow_lite.reducestreamer import Datum, Message, Metadata, ReduceStreamAsyncServer, ReduceStreamer -class ReduceCounter(reducestreamer.ReduceStreamer): +class ReduceCounter(ReduceStreamer): """ A reduce streaming counter that emits intermediate results. This demonstrates the key difference from regular Reducer: - - Regular Reducer: waits for all data, then returns Messages + - Regular Reducer: waits for all data, then returns list[Message] - ReduceStreamer: yields Message objects incrementally as an async iterator """ def __init__(self, initial: int = 0) -> None: self.counter = initial - async def handler( - self, - keys: list[str], - datums: AsyncIterable[reducestreamer.Datum], - md: reducestreamer.Metadata, - ) -> AsyncIterator[reducestreamer.Message]: + async def handler(self, keys: list[str], datums: AsyncIterator[Datum], md: Metadata) -> AsyncIterator[Message]: """ Process datums and yield messages incrementally. Args: keys: List of keys for this window - datums: Async iterable of incoming data + datums: Async iterator of incoming data md: Metadata containing window information Yields: @@ -52,42 +47,27 @@ async def handler( # Emit intermediate result every 10 items if self.counter % 10 == 0: - msg = ( - f"counter:{self.counter} interval_window_start:{iw.start} interval_window_end:{iw.end}" - ).encode() + msg = f"counter:{self.counter} interval_window_start:{iw.start} interval_window_end:{iw.end}" print(f"Yielding intermediate result: counter={self.counter}") # Early release of data - this is the key feature of reduce streaming! - yield reducestreamer.Message(msg, keys=keys) + yield Message(msg.encode(), keys=keys) # Emit final result - msg = ( - f"counter:{self.counter} (FINAL) interval_window_start:{iw.start} interval_window_end:{iw.end}" - ).encode() + msg = f"counter:{self.counter} (FINAL) interval_window_start:{iw.start} interval_window_end:{iw.end}" print(f"Yielding final result: counter={self.counter}") - yield reducestreamer.Message(msg, keys=keys) + yield Message(msg.encode(), keys=keys) -async def start(creator: type, init_args: tuple): - """Start the reduce stream server.""" - sock_file = "/var/run/numaflow/reducestream.sock" - server_info_file = "/var/run/numaflow/reducestreamer-server-info" - server = reducestreamer.ReduceStreamAsyncServer(sock_file, server_info_file) - - 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: - print("Starting Reduce Stream Counter Server...") - await server.start(creator, init_args) - print("Shutting down gracefully...") - except asyncio.CancelledError: - server.stop() - return +async def main() -> None: + # ReduceStreamAsyncServer needs a factory: a callable that takes no arguments + # and returns a new ReduceStreamer. The server calls it at the start of every + # window (once for each set of keys), so each window starts with its own state. + # functools.partial binds initial=0, so the server can call + # ReduceCounter(initial=0) without arguments. + # `serve` returns when SIGINT or SIGTERM arrives. + print("Starting Reduce Stream Counter Server...") + await ReduceStreamAsyncServer(functools.partial(ReduceCounter, initial=0)).serve() if __name__ == "__main__": - asyncio.run(start(ReduceCounter, (0,))) + asyncio.run(main()) diff --git a/packages/pynumaflow-lite/manifests/reducestream/uv.lock b/packages/pynumaflow-lite/manifests/reducestream/uv.lock new file mode 100644 index 00000000..b951ff54 --- /dev/null +++ b/packages/pynumaflow-lite/manifests/reducestream/uv.lock @@ -0,0 +1,8 @@ +version = 1 +revision = 3 +requires-python = "==3.11.*" + +[[package]] +name = "reducestream-counter" +version = "0.1.0" +source = { virtual = "." } diff --git a/packages/pynumaflow-lite/pynumaflow_lite/__init__.py b/packages/pynumaflow-lite/pynumaflow_lite/__init__.py index 2594994b..80e5ed1b 100644 --- a/packages/pynumaflow-lite/pynumaflow_lite/__init__.py +++ b/packages/pynumaflow-lite/pynumaflow_lite/__init__.py @@ -13,6 +13,7 @@ from ._mapstream_server import MapStreamAsyncServer from ._reduce_dtypes import Reducer from ._reduce_server import ReduceAsyncServer +from ._reducestream_server import ReduceStreamAsyncServer from ._reducestreamer_dtypes import ReduceStreamer from ._session_reduce_dtypes import SessionReducer from ._sideinput_dtypes import SideInput @@ -49,6 +50,7 @@ reducer.ReduceAsyncServer = ReduceAsyncServer session_reducer.SessionReducer = SessionReducer reducestreamer.ReduceStreamer = ReduceStreamer +reducestreamer.ReduceStreamAsyncServer = ReduceStreamAsyncServer accumulator.Accumulator = Accumulator sinker.Sinker = Sinker sinker.SinkAsyncServer = SinkAsyncServer diff --git a/packages/pynumaflow-lite/pynumaflow_lite/_batchmap_server.py b/packages/pynumaflow-lite/pynumaflow_lite/_batchmap_server.py index da2566c2..10dcae84 100644 --- a/packages/pynumaflow-lite/pynumaflow_lite/_batchmap_server.py +++ b/packages/pynumaflow-lite/pynumaflow_lite/_batchmap_server.py @@ -1,6 +1,6 @@ from __future__ import annotations -from collections.abc import AsyncIterable, Awaitable, Callable +from collections.abc import AsyncIterator, Awaitable, Callable from typing import TypeAlias from ._async_server import _AsyncServer @@ -15,7 +15,7 @@ class BatchMapAsyncServer(_AsyncServer): def __init__( self, - handler: Callable[[AsyncIterable[Datum]], Awaitable[list[BatchResponse]]], + handler: Callable[[AsyncIterator[Datum]], Awaitable[list[BatchResponse]]], *, sock_file: str | None = None, server_info_file: str | None = None, diff --git a/packages/pynumaflow-lite/pynumaflow_lite/_batchmapper_dtypes.py b/packages/pynumaflow-lite/pynumaflow_lite/_batchmapper_dtypes.py index 557c8afe..effbd7a8 100644 --- a/packages/pynumaflow-lite/pynumaflow_lite/_batchmapper_dtypes.py +++ b/packages/pynumaflow-lite/pynumaflow_lite/_batchmapper_dtypes.py @@ -1,5 +1,5 @@ from abc import ABCMeta, abstractmethod -from collections.abc import AsyncIterable +from collections.abc import AsyncIterator from pynumaflow_lite.batchmapper import BatchResponse, Datum @@ -13,7 +13,7 @@ def __call__(self, *args, **kwargs): return self.handler(*args, **kwargs) @abstractmethod - async def handler(self, batch: AsyncIterable[Datum]) -> list[BatchResponse]: + async def handler(self, batch: AsyncIterator[Datum]) -> list[BatchResponse]: """ Implement this handler function for batch mapping. The returned list length should equal the input batch size. diff --git a/packages/pynumaflow-lite/pynumaflow_lite/_reduce_dtypes.py b/packages/pynumaflow-lite/pynumaflow_lite/_reduce_dtypes.py index 947c9fc5..c4fc00b2 100644 --- a/packages/pynumaflow-lite/pynumaflow_lite/_reduce_dtypes.py +++ b/packages/pynumaflow-lite/pynumaflow_lite/_reduce_dtypes.py @@ -1,5 +1,5 @@ from abc import ABCMeta, abstractmethod -from collections.abc import AsyncIterable +from collections.abc import AsyncIterator from pynumaflow_lite.reducer import Datum, Message, Metadata @@ -13,8 +13,8 @@ def __call__(self, *args, **kwargs): return self.handler(*args, **kwargs) @abstractmethod - async def handler(self, keys: list[str], datums: AsyncIterable[Datum], md: Metadata) -> list[Message]: + async def handler(self, keys: list[str], datums: AsyncIterator[Datum], md: Metadata) -> list[Message]: """ - Implement this handler; consume `datums` async iterable and return the messages. + Implement this handler; consume `datums` async iterator and return the messages. """ pass diff --git a/packages/pynumaflow-lite/pynumaflow_lite/_reduce_server.py b/packages/pynumaflow-lite/pynumaflow_lite/_reduce_server.py index 83c21406..6489323e 100644 --- a/packages/pynumaflow-lite/pynumaflow_lite/_reduce_server.py +++ b/packages/pynumaflow-lite/pynumaflow_lite/_reduce_server.py @@ -1,7 +1,7 @@ from __future__ import annotations import inspect -from collections.abc import AsyncIterable, Awaitable, Callable +from collections.abc import AsyncIterator, Awaitable, Callable from typing import TypeAlias from ._async_server import _AsyncServer @@ -12,7 +12,7 @@ Message: TypeAlias = _reducer.Message Metadata: TypeAlias = _reducer.Metadata -ReduceHandler: TypeAlias = Callable[[list[str], AsyncIterable[Datum], Metadata], Awaitable[list[Message]]] +ReduceHandler: TypeAlias = Callable[[list[str], AsyncIterator[Datum], Metadata], Awaitable[list[Message]]] class ReduceAsyncServer(_AsyncServer): diff --git a/packages/pynumaflow-lite/pynumaflow_lite/_reducestream_server.py b/packages/pynumaflow-lite/pynumaflow_lite/_reducestream_server.py new file mode 100644 index 00000000..ca95a5b9 --- /dev/null +++ b/packages/pynumaflow-lite/pynumaflow_lite/_reducestream_server.py @@ -0,0 +1,47 @@ +from __future__ import annotations + +import inspect +from collections.abc import AsyncIterator, Callable +from typing import TypeAlias + +from ._async_server import _AsyncServer +from ._reducestreamer_dtypes import ReduceStreamer +from .pynumaflow_lite import reducestreamer as _reducestreamer + +Datum: TypeAlias = _reducestreamer.Datum +Message: TypeAlias = _reducestreamer.Message +Metadata: TypeAlias = _reducestreamer.Metadata + +ReduceStreamHandler: TypeAlias = Callable[[list[str], AsyncIterator[Datum], Metadata], AsyncIterator[Message]] + + +class ReduceStreamAsyncServer(_AsyncServer): + """Reduce stream server. + + `handler` is an async generator function, or a zero-argument factory that + returns a handler, such as a `ReduceStreamer` subclass or + `functools.partial(MyReduceStreamer, x=1)`. The server calls a factory once + for each window and key set, so each call gets its own state. + """ + + _name = "reducestream" + + def __init__( + self, + handler: ReduceStreamHandler | Callable[[], ReduceStreamHandler], + *, + sock_file: str | None = None, + server_info_file: str | None = None, + install_signal_handlers: bool = True, + ) -> None: + if isinstance(handler, ReduceStreamer): + raise TypeError( + "pass the ReduceStreamer class or a factory, not an instance: " + "the server creates a new instance for each window" + ) + factory = (lambda: handler) if inspect.isasyncgenfunction(handler) else handler + super().__init__( + _reducestreamer._ReduceStreamAsyncServer(sock_file, server_info_file), + factory, + install_signal_handlers=install_signal_handlers, + ) diff --git a/packages/pynumaflow-lite/pynumaflow_lite/_reducestreamer_dtypes.py b/packages/pynumaflow-lite/pynumaflow_lite/_reducestreamer_dtypes.py index 25c65f94..ca1cb18b 100644 --- a/packages/pynumaflow-lite/pynumaflow_lite/_reducestreamer_dtypes.py +++ b/packages/pynumaflow-lite/pynumaflow_lite/_reducestreamer_dtypes.py @@ -1,5 +1,5 @@ from abc import ABCMeta, abstractmethod -from collections.abc import AsyncIterable, AsyncIterator +from collections.abc import AsyncIterator from pynumaflow_lite.reducestreamer import Datum, Message, Metadata @@ -16,9 +16,9 @@ def __call__(self, *args, **kwargs): return self.handler(*args, **kwargs) @abstractmethod - async def handler(self, keys: list[str], datums: AsyncIterable[Datum], md: Metadata) -> AsyncIterator[Message]: + async def handler(self, keys: list[str], datums: AsyncIterator[Datum], md: Metadata) -> AsyncIterator[Message]: """ - Implement this handler; consume `datums` async iterable and yield Messages incrementally. + Implement this handler; consume `datums` async iterator and yield Messages incrementally. Args: keys: List of keys for this window @@ -28,4 +28,5 @@ async def handler(self, keys: list[str], datums: AsyncIterable[Datum], md: Metad Yields: Message objects to be sent to the next vertex """ - pass + raise NotImplementedError + yield # makes this an async generator, so overrides that yield type-check diff --git a/packages/pynumaflow-lite/pynumaflow_lite/batchmapper.pyi b/packages/pynumaflow-lite/pynumaflow_lite/batchmapper.pyi index fde5a076..1ff23749 100644 --- a/packages/pynumaflow-lite/pynumaflow_lite/batchmapper.pyi +++ b/packages/pynumaflow-lite/pynumaflow_lite/batchmapper.pyi @@ -1,5 +1,5 @@ import datetime as _dt -from collections.abc import AsyncIterable, Awaitable, Callable +from collections.abc import AsyncIterator, Awaitable, Callable from types import TracebackType from typing import Self @@ -78,14 +78,14 @@ class _BatchMapAsyncServer: sock_file: str | None = ..., server_info_file: str | None = ..., ) -> None: ... - def start(self, handler: Callable[[AsyncIterable[Datum]], Awaitable[list[BatchResponse]]]) -> Awaitable[None]: ... + def start(self, handler: Callable[[AsyncIterator[Datum]], Awaitable[list[BatchResponse]]]) -> Awaitable[None]: ... def wait_ready(self) -> Awaitable[None]: ... def stop(self) -> None: ... class BatchMapAsyncServer: def __init__( self, - handler: Callable[[AsyncIterable[Datum]], Awaitable[list[BatchResponse]]], + handler: Callable[[AsyncIterator[Datum]], Awaitable[list[BatchResponse]]], *, sock_file: str | None = ..., server_info_file: str | None = ..., diff --git a/packages/pynumaflow-lite/pynumaflow_lite/reducer.pyi b/packages/pynumaflow-lite/pynumaflow_lite/reducer.pyi index 87228882..13976cd6 100644 --- a/packages/pynumaflow-lite/pynumaflow_lite/reducer.pyi +++ b/packages/pynumaflow-lite/pynumaflow_lite/reducer.pyi @@ -1,5 +1,5 @@ import datetime as _dt -from collections.abc import AsyncIterable, Awaitable, Callable +from collections.abc import AsyncIterator, Awaitable, Callable from types import TracebackType from typing import Self, TypeAlias @@ -52,7 +52,7 @@ class Metadata: def __init__(self, *, interval_window: IntervalWindow) -> None: ... -_ReduceHandler: TypeAlias = Callable[[list[str], AsyncIterable[Datum], Metadata], Awaitable[list[Message]]] +_ReduceHandler: TypeAlias = Callable[[list[str], AsyncIterator[Datum], Metadata], Awaitable[list[Message]]] class _ReduceAsyncServer: def __init__( diff --git a/packages/pynumaflow-lite/pynumaflow_lite/reducestreamer.pyi b/packages/pynumaflow-lite/pynumaflow_lite/reducestreamer.pyi index 92c3244f..44108dc5 100644 --- a/packages/pynumaflow-lite/pynumaflow_lite/reducestreamer.pyi +++ b/packages/pynumaflow-lite/pynumaflow_lite/reducestreamer.pyi @@ -1,5 +1,7 @@ import datetime as _dt -from collections.abc import AsyncIterable, AsyncIterator, Awaitable, Callable +from collections.abc import AsyncIterator, Awaitable, Callable +from types import TracebackType +from typing import Self, TypeAlias # Re-export the Python ABC for user convenience and typing from ._reducestreamer_dtypes import ReduceStreamer as ReduceStreamer @@ -17,6 +19,8 @@ class Message: ) -> None: ... @staticmethod def to_drop() -> Message: ... + def __repr__(self) -> str: ... + def __eq__(self, other: object) -> bool: ... class Datum: keys: list[str] @@ -25,6 +29,15 @@ class Datum: 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: ... @@ -32,39 +45,54 @@ class IntervalWindow: start: _dt.datetime end: _dt.datetime + def __init__(self, start: _dt.datetime, end: _dt.datetime) -> None: ... + class Metadata: interval_window: IntervalWindow -class PyAsyncDatumStream: - """ - Python-visible async iterator that yields Datum items from a Tokio mpsc channel. - """ + def __init__(self, *, interval_window: IntervalWindow) -> None: ... - def __init__(self) -> None: ... - def __aiter__(self) -> PyAsyncDatumStream: ... - def __anext__(self) -> Datum: ... +_ReduceStreamHandler: TypeAlias = Callable[[list[str], AsyncIterator[Datum], Metadata], AsyncIterator[Message]] -class ReduceStreamAsyncServer: +class _ReduceStreamAsyncServer: def __init__( self, - sock_file: str = ..., - info_file: str = ..., + sock_file: str | None = ..., + server_info_file: str | None = ..., ) -> None: ... def start( self, - py_creator: ( - type[ReduceStreamer] | Callable[[list[str], AsyncIterable[Datum], Metadata], AsyncIterator[Message]] - ), - init_args: tuple | None = ..., + factory: Callable[[], _ReduceStreamHandler], ) -> Awaitable[None]: ... + def wait_ready(self) -> Awaitable[None]: ... + def stop(self) -> None: ... + +class ReduceStreamAsyncServer: + def __init__( + self, + handler: _ReduceStreamHandler | Callable[[], _ReduceStreamHandler], + *, + sock_file: str | None = ..., + server_info_file: str | None = ..., + install_signal_handlers: bool = ..., + ) -> None: ... + async def serve(self) -> None: ... def stop(self) -> None: ... + async def wait_ready(self) -> None: ... + async def wait_for_termination(self) -> None: ... + async def __aenter__(self) -> Self: ... + async def __aexit__( + self, + exc_type: type[BaseException] | None, + exc: BaseException | None, + tb: TracebackType | None, + ) -> None: ... __all__ = [ "Datum", "IntervalWindow", "Message", "Metadata", - "PyAsyncDatumStream", "ReduceStreamAsyncServer", "ReduceStreamer", ] diff --git a/packages/pynumaflow-lite/src/batchmap/server.rs b/packages/pynumaflow-lite/src/batchmap/server.rs index ebc08a32..1685091d 100644 --- a/packages/pynumaflow-lite/src/batchmap/server.rs +++ b/packages/pynumaflow-lite/src/batchmap/server.rs @@ -47,7 +47,7 @@ impl batchmap::BatchMapper for PyBatchMapRunner { // When input ends, dropping tx closes the channel }); - // Call the Python coroutine: py_func(batch: AsyncIterable[Datum]) -> list[BatchResponse] + // Call the Python coroutine: py_func(batch: AsyncIterator[Datum]) -> list[BatchResponse] let fut = Python::attach(|py| -> PyResult<_> { let locals = pyo3_async_runtimes::TaskLocals::new(self.event_loop.bind(py).clone()); let py_func = self.py_func.clone(); diff --git a/packages/pynumaflow-lite/src/reduce/mod.rs b/packages/pynumaflow-lite/src/reduce/mod.rs index 38715b85..d89110ed 100644 --- a/packages/pynumaflow-lite/src/reduce/mod.rs +++ b/packages/pynumaflow-lite/src/reduce/mod.rs @@ -192,14 +192,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 } @@ -237,8 +229,9 @@ async fn reduce_client(sock_file: String) -> PyResult list[Message] + // factory()(keys: list[str], datums: AsyncIterator[Datum], md: Metadata) -> list[Message] let fut = match Python::attach(|py| -> PyResult<_> { let locals = pyo3_async_runtimes::TaskLocals::new(self.event_loop.bind(py).clone()); let handler = self.factory.call0(py)?; diff --git a/packages/pynumaflow-lite/src/reducestream/mod.rs b/packages/pynumaflow-lite/src/reducestream/mod.rs index be9c7816..9127e043 100644 --- a/packages/pynumaflow-lite/src/reducestream/mod.rs +++ b/packages/pynumaflow-lite/src/reducestream/mod.rs @@ -1,5 +1,7 @@ use std::sync::Mutex; +use numaflow::reducestream; + pub mod server; use pyo3::prelude::*; @@ -7,53 +9,64 @@ use pyo3::prelude::*; // Re-export types from reduce module that are shared pub use crate::reduce::{Datum, IntervalWindow, Message, Metadata, PyAsyncDatumStream}; -/// Async Reduce Stream Server that can be started from Python code, taking a class (creator). -#[pyclass(module = "pynumaflow_lite.reducestreamer")] +/// Async Reduce Stream Server that can be started from Python code, taking a reducer factory. +#[pyclass( + name = "_ReduceStreamAsyncServer", + module = "pynumaflow_lite.reducestreamer" +)] pub struct ReduceStreamAsyncServer { sock_file: String, - info_file: String, + server_info_file: String, shutdown_tx: Mutex>>, } #[pymethods] impl ReduceStreamAsyncServer { #[new] - #[pyo3(signature = (sock_file="/var/run/numaflow/reducestream.sock".to_string(), info_file="/var/run/numaflow/reducestreamer-server-info".to_string()))] - fn new(sock_file: String, info_file: String) -> Self { + #[pyo3(signature = ( + sock_file: "str | None"=None, + server_info_file: "str | None"=None, + ) -> "_ReduceStreamAsyncServer")] + fn new(sock_file: Option, server_info_file: Option) -> Self { Self { - sock_file, - info_file, + sock_file: sock_file.unwrap_or_else(|| reducestream::SOCK_ADDR.to_string()), + server_info_file: server_info_file + .unwrap_or_else(|| reducestream::SERVER_INFO_FILE.to_string()), shutdown_tx: Mutex::new(None), } } - /// Start the server with the given Python class (creator). - /// Only class-based reducers are supported (not function-based). - /// init_args is an optional tuple of positional arguments to pass to the class constructor. - #[pyo3(signature = (py_creator, init_args=None))] - pub fn start<'a>( - &self, - py: Python<'a>, - py_creator: Py, - init_args: Option>, - ) -> PyResult> { + /// Start the server with the given zero-argument factory. The server calls the + /// factory once for each window and key set to get the async generator handler. + #[pyo3(signature = (factory: "callable") -> "None")] + pub fn start<'a>(&self, py: Python<'a>, factory: 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(); *guard = Some(tx); } + pyo3_async_runtimes::tokio::future_into_py( + py, + crate::reducestream::server::start(factory, sock_file, server_info_file, rx), + ) + } + + /// Wait until the Numaflow IsReady probe succeeds over the reducestream UDS. + #[pyo3(signature = () -> "None")] + pub fn wait_ready<'a>(&self, py: Python<'a>) -> PyResult> { + let sock_file = self.sock_file.clone(); + pyo3_async_runtimes::tokio::future_into_py(py, async move { - crate::reducestream::server::start(py_creator, init_args, sock_file, info_file, rx) - .await - .expect("server failed to start"); + crate::reduce::wait_for_ready(sock_file).await; Ok(()) }) } /// Trigger server shutdown from Python (idempotent). + #[pyo3(signature = () -> "None")] pub fn stop(&self) -> PyResult<()> { if let Some(tx) = self.shutdown_tx.lock().unwrap().take() { let _ = tx.send(()); @@ -69,6 +82,5 @@ pub(crate) fn populate_py_module(m: &Bound) -> PyResult<()> { m.add_class::()?; m.add_class::()?; m.add_class::()?; - m.add_class::()?; Ok(()) } diff --git a/packages/pynumaflow-lite/src/reducestream/server.rs b/packages/pynumaflow-lite/src/reducestream/server.rs index 3d11b23f..149be39a 100644 --- a/packages/pynumaflow-lite/src/reducestream/server.rs +++ b/packages/pynumaflow-lite/src/reducestream/server.rs @@ -1,55 +1,61 @@ +use std::sync::{Arc, Mutex}; + use crate::pyiterables::PyAsyncIterStream; +use crate::pyrs::{combine_errors, format_error}; use crate::reducestream::{ Datum as PyDatum, IntervalWindow as PyIntervalWindow, Message as PyMessage, Metadata as PyMetadata, PyAsyncDatumStream, }; use numaflow::reducestream; use numaflow::shared::ServerExtras; +use pyo3::exceptions::PyTypeError; use pyo3::prelude::*; -use std::sync::Arc; use tokio_stream::StreamExt; pub(crate) struct PyReduceStreamerCreator { /// handle to Python event loop pub(crate) event_loop: Arc>, - /// Python class to instantiate per window - pub(crate) py_creator: Arc>, - /// optional tuple of positional args - pub(crate) init_args: Option>>, + /// Zero-argument Python callable that returns the handler for one window + pub(crate) factory: Arc>, + pub(crate) errors: Arc>>, } pub(crate) struct PyReduceStreamerRunner { /// handle to Python event loop pub(crate) event_loop: Arc>, - /// Instance of class per window - pub(crate) py_instance: Arc>, + pub(crate) factory: Arc>, + pub(crate) errors: Arc>>, } impl reducestream::ReduceStreamerCreator for PyReduceStreamerCreator { type R = PyReduceStreamerRunner; + // The factory runs later, in reducestream(). A panic in create() would stop the + // numaflow task executor without an error to the client, but reducestream() + // reports a panic. fn create(&self) -> Self::R { - // Instantiate the Python class synchronously under the GIL. - let inst = Python::attach(|py| { - let class = self.py_creator.as_ref(); - match &self.init_args { - Some(args) => { - let bound = args.as_ref().bind(py); - let py_tuple = bound.cast()?; - class.call1(py, py_tuple) - } - None => class.call0(py), - } - }) - .expect("failed to instantiate Python reduce streamer class"); - PyReduceStreamerRunner { event_loop: self.event_loop.clone(), - py_instance: Arc::new(inst), + factory: self.factory.clone(), + errors: self.errors.clone(), } } } +impl PyReduceStreamerRunner { + fn fail(&self, error: PyErr) -> ! { + // numaflow calls reducestream() concurrently, so each error belongs to a + // different window. Keep all of them. start() raises them together, which + // lets Python format every traceback instead of Rust printing them by hand. + let message = Python::attach(|py| format_error(py, &error)); + self.errors.lock().unwrap().push(error); + + // numaflow catches this panic, sends a gRPC error for this window, and + // starts the server shutdown. + panic!("{message}"); + } +} + #[tonic::async_trait] impl reducestream::ReduceStreamer for PyReduceStreamerRunner { async fn reducestream( @@ -65,8 +71,7 @@ impl reducestream::ReduceStreamer for PyReduceStreamerRunner { // Spawn a task forwarding incoming datums to the Python-facing channel let forwarder = tokio::spawn(async move { while let Some(req) = input.recv().await { - let datum = PyDatum::from(req); - if tx.send(datum).await.is_err() { + if tx.send(req.into()).await.is_err() { break; } } @@ -74,44 +79,58 @@ impl reducestream::ReduceStreamer for PyReduceStreamerRunner { // the Python async iterable will stop. }); - // Call the Python coroutine: - // obj.handler(keys, datums: AsyncIterable[Datum], md: Metadata) -> AsyncIterator[Message] - let agen_obj = Python::attach(|py| { - let obj = self.py_instance.clone(); - let stream = PyAsyncDatumStream::new_with(rx); + // Call the Python handler: + // factory()(keys: list[str], datums: AsyncIterator[Datum], md: Metadata) -> AsyncIterator[Message] + let agen = match Python::attach(|py| -> PyResult<_> { + let handler = self.factory.call0(py)?; - // Build metadata object + let stream = PyAsyncDatumStream::new_with(rx); let interval = PyIntervalWindow::new(md.interval_window.start_time, md.interval_window.end_time); let md_py = PyMetadata::new(interval); - // Call handler method - let agen = obj - .call_method1(py, "handler", (keys, stream, md_py)) - .expect("python handler method raised before returning async iterable"); - - // Keep as Py - agen.clone_ref(py).extract(py).unwrap_or(agen) - }); - - // Wrap the Python AsyncIterable in a Rust Stream that yields incrementally - let mut stream = PyAsyncIterStream::::new(agen_obj, self.event_loop.clone()) - .expect("failed to construct PyAsyncIterStream"); + handler.call1(py, (keys, stream, md_py)) + }) { + Ok(agen) => agen, + Err(error) => self.fail(error), + }; + + // Wrap the Python AsyncIterable in a Rust Stream that yields incrementally. + // Items stay as Py so that a wrong item type gives a clear error below. + let mut stream = match PyAsyncIterStream::>::new(agen, self.event_loop.clone()) { + Ok(stream) => stream, + Err(_) => self.fail(PyErr::new::( + "reducestream handler must be an async generator (return AsyncIterator[Message])", + )), + }; // Forward each yielded message immediately to the output sender while let Some(item) = stream.next().await { - match item { - Ok(py_msg) => { - let out: reducestream::Message = py_msg.into(); - if output.send(out).await.is_err() { - break; - } - } - Err(e) => { - // Non-stop errors are surfaced per-item; log and stop this stream. - eprintln!("Python async iteration error in reduce streamer: {:?}", e); - break; - } + let obj = match item { + Ok(obj) => obj, + Err(error) => self.fail(error), + }; + + let message: PyMessage = match Python::attach(|py| { + obj.extract(py).map_err(|_| { + let type_name = obj + .bind(py) + .get_type() + .name() + .map(|name| name.to_string_lossy().into_owned()) + .unwrap_or_else(|_| "".to_string()); + PyErr::new::(format!( + "reducestream handler must yield Message, got {type_name}" + )) + }) + }) { + Ok(message) => message, + Err(error) => self.fail(error), + }; + + // The receiver is gone, so the client does not want more messages. + if output.send(message.into()).await.is_err() { + break; } } @@ -122,22 +141,30 @@ impl reducestream::ReduceStreamer for PyReduceStreamerRunner { /// Start the reduce stream server by spinning up a dedicated Python asyncio loop and wiring shutdown. pub(super) async fn start( - py_creator: Py, - init_args: Option>, + factory: Py, sock_file: String, info_file: String, shutdown_rx: tokio::sync::oneshot::Receiver<()>, ) -> Result<(), pyo3::PyErr> { let (tx, rx) = tokio::sync::oneshot::channel(); - let py_asyncio_loop_handle = tokio::task::spawn_blocking(move || crate::pyrs::run_asyncio(tx)); + let py_asyncio_loop_handle = tokio::task::spawn_blocking({ + println!( + "Starting ReduceStream 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 errors = Arc::new(Mutex::new(Vec::new())); + // Shutdown has two sources, and neither one needs a channel here. The Python + // side signals stop() through shutdown_rx. An uncaught Python error panics in + // fail(), and numaflow then shuts the server down on its own. let creator = PyReduceStreamerCreator { event_loop: event_loop.clone(), - py_creator: Arc::new(py_creator), - init_args: init_args.map(Arc::new), + factory: Arc::new(factory), + errors: Arc::clone(&errors), }; let server = reducestream::Server::new(creator) @@ -145,7 +172,7 @@ pub(super) async fn start( .with_server_info_file(info_file); let result = server - .start_with_shutdown(combined_rx) + .start_with_shutdown(shutdown_rx) .await .map_err(|e| pyo3::PyErr::new::(e.to_string())); @@ -156,15 +183,14 @@ pub(super) async fn start( } }); - println!("Numaflow Core (reduce stream) has shutdown..."); + println!("Numaflow ReduceStream has shutdown..."); // Wait for the blocking asyncio thread to finish. let _ = py_asyncio_loop_handle.await; - // if not finished, abort it - if !sig_handle.is_finished() { - println!("Aborting signal handler"); - sig_handle.abort(); + let errors = std::mem::take(&mut *errors.lock().unwrap()); + if !errors.is_empty() { + return Err(Python::attach(|py| combine_errors(py, errors))); } result diff --git a/packages/pynumaflow-lite/tests/examples/batchmap_cat.py b/packages/pynumaflow-lite/tests/examples/batchmap_cat.py index 5fd343b1..c48b1402 100644 --- a/packages/pynumaflow-lite/tests/examples/batchmap_cat.py +++ b/packages/pynumaflow-lite/tests/examples/batchmap_cat.py @@ -1,11 +1,11 @@ import asyncio -from collections.abc import AsyncIterable +from collections.abc import AsyncIterator from pynumaflow_lite.batchmapper import BatchMapAsyncServer, BatchResponse, Datum, Message async def async_handler( - batch: AsyncIterable[Datum], + batch: AsyncIterator[Datum], ) -> list[BatchResponse]: responses = [] async for d in batch: diff --git a/packages/pynumaflow-lite/tests/examples/batchmap_cat_class.py b/packages/pynumaflow-lite/tests/examples/batchmap_cat_class.py index 36dbacf8..8b138aed 100644 --- a/packages/pynumaflow-lite/tests/examples/batchmap_cat_class.py +++ b/packages/pynumaflow-lite/tests/examples/batchmap_cat_class.py @@ -1,11 +1,11 @@ import asyncio -from collections.abc import AsyncIterable +from collections.abc import AsyncIterator from pynumaflow_lite.batchmapper import BatchMapAsyncServer, BatchMapper, BatchResponse, Datum, Message class SimpleBatchCat(BatchMapper): - async def handler(self, batch: AsyncIterable[Datum]) -> list[BatchResponse]: + async def handler(self, batch: AsyncIterator[Datum]) -> list[BatchResponse]: responses = [] async for d in batch: if d.value == b"bad world": diff --git a/packages/pynumaflow-lite/tests/examples/reduce_counter_class.py b/packages/pynumaflow-lite/tests/examples/reduce_counter_class.py index 92b632be..136e830f 100644 --- a/packages/pynumaflow-lite/tests/examples/reduce_counter_class.py +++ b/packages/pynumaflow-lite/tests/examples/reduce_counter_class.py @@ -1,6 +1,6 @@ import asyncio import functools -from collections.abc import AsyncIterable +from collections.abc import AsyncIterator from pynumaflow_lite.reducer import Datum, Message, Metadata, ReduceAsyncServer, Reducer @@ -9,7 +9,7 @@ class ReduceCounter(Reducer): def __init__(self, initial: int = 0) -> None: self.initial = initial - async def handler(self, keys: list[str], datums: AsyncIterable[Datum], md: Metadata) -> list[Message]: + async def handler(self, keys: list[str], datums: AsyncIterator[Datum], md: Metadata) -> list[Message]: counter = self.initial async for _ in datums: counter += 1 diff --git a/packages/pynumaflow-lite/tests/examples/reduce_counter_func.py b/packages/pynumaflow-lite/tests/examples/reduce_counter_func.py index af168d6b..ee52efa4 100644 --- a/packages/pynumaflow-lite/tests/examples/reduce_counter_func.py +++ b/packages/pynumaflow-lite/tests/examples/reduce_counter_func.py @@ -1,10 +1,10 @@ import asyncio -from collections.abc import AsyncIterable +from collections.abc import AsyncIterator from pynumaflow_lite.reducer import Datum, Message, Metadata, ReduceAsyncServer -async def reduce_handler(keys: list[str], datums: AsyncIterable[Datum], md: Metadata) -> list[Message]: +async def reduce_handler(keys: list[str], datums: AsyncIterator[Datum], md: Metadata) -> list[Message]: interval_window = md.interval_window counter = 0 async for _ in datums: diff --git a/packages/pynumaflow-lite/tests/examples/reducestream_counter.py b/packages/pynumaflow-lite/tests/examples/reducestream_counter.py index 5988bc86..a723d152 100644 --- a/packages/pynumaflow-lite/tests/examples/reducestream_counter.py +++ b/packages/pynumaflow-lite/tests/examples/reducestream_counter.py @@ -1,11 +1,11 @@ import asyncio -import signal -from collections.abc import AsyncIterable, AsyncIterator +import functools +from collections.abc import AsyncIterator -from pynumaflow_lite import reducestreamer +from pynumaflow_lite.reducestreamer import Datum, Message, Metadata, ReduceStreamAsyncServer, ReduceStreamer -class ReduceStreamCounter(reducestreamer.ReduceStreamer): +class ReduceStreamCounter(ReduceStreamer): """ A reduce streaming counter that emits intermediate results. @@ -16,47 +16,27 @@ class ReduceStreamCounter(reducestreamer.ReduceStreamer): def __init__(self, initial: int = 0) -> None: self.counter = initial - async def handler( - self, - keys: list[str], - datums: AsyncIterable[reducestreamer.Datum], - md: reducestreamer.Metadata, - ) -> AsyncIterator[reducestreamer.Message]: + async def handler(self, keys: list[str], datums: AsyncIterator[Datum], md: Metadata) -> AsyncIterator[Message]: """ Process datums and yield messages incrementally. For testing purposes, we yield a message for each datum received. """ iw = md.interval_window - self.counter = 0 - async for _ in datums: self.counter += 1 # Yield a message for each datum (streaming behavior) - msg = (f"counter:{self.counter} interval_window_start:{iw.start} interval_window_end:{iw.end}").encode() - yield reducestreamer.Message(msg, keys=keys) - - -async def start(creator: type, init_args: tuple): - """Start the reduce stream server.""" - sock_file = "/tmp/var/run/numaflow/reducestream.sock" - server_info_file = "/tmp/var/run/numaflow/reducestreamer-server-info" - server = reducestreamer.ReduceStreamAsyncServer(sock_file, server_info_file) + msg = f"counter:{self.counter} interval_window_start:{iw.start} interval_window_end:{iw.end}" + yield Message(msg.encode(), keys=keys) - 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(creator, init_args) - print("Shutting down gracefully...") - except asyncio.CancelledError: - server.stop() - return +async def main() -> None: + await ReduceStreamAsyncServer( + functools.partial(ReduceStreamCounter, initial=0), + sock_file="/tmp/var/run/numaflow/reducestream.sock", + server_info_file="/tmp/var/run/numaflow/reducestreamer-server-info", + ).serve() if __name__ == "__main__": - asyncio.run(start(ReduceStreamCounter, (0,))) + asyncio.run(main()) diff --git a/packages/pynumaflow-lite/tests/examples/reducestream_counter_func.py b/packages/pynumaflow-lite/tests/examples/reducestream_counter_func.py new file mode 100644 index 00000000..d0136221 --- /dev/null +++ b/packages/pynumaflow-lite/tests/examples/reducestream_counter_func.py @@ -0,0 +1,25 @@ +import asyncio +from collections.abc import AsyncIterator + +from pynumaflow_lite.reducestreamer import Datum, Message, Metadata, ReduceStreamAsyncServer + + +async def reducestream_handler(keys: list[str], datums: AsyncIterator[Datum], md: Metadata) -> AsyncIterator[Message]: + iw = md.interval_window + counter = 0 + async for _ in datums: + counter += 1 + msg = f"counter:{counter} interval_window_start:{iw.start} interval_window_end:{iw.end}" + yield Message(msg.encode(), keys=keys) + + +async def main() -> None: + await ReduceStreamAsyncServer( + reducestream_handler, + sock_file="/tmp/var/run/numaflow/reducestream.sock", + server_info_file="/tmp/var/run/numaflow/reducestreamer-server-info", + ).serve() + + +if __name__ == "__main__": + asyncio.run(main()) diff --git a/packages/pynumaflow-lite/tests/test_reduce_unit.py b/packages/pynumaflow-lite/tests/test_reduce_unit.py index b0a6b45f..a4b5ce70 100644 --- a/packages/pynumaflow-lite/tests/test_reduce_unit.py +++ b/packages/pynumaflow-lite/tests/test_reduce_unit.py @@ -3,7 +3,7 @@ import os import subprocess import uuid -from collections.abc import AsyncIterable, AsyncIterator +from collections.abc import AsyncIterator from contextlib import suppress from datetime import UTC, datetime from pathlib import Path @@ -22,7 +22,7 @@ def __init__(self, initial: int = 0) -> None: self.counter = initial async def handler( - self, keys: list[str], datums: AsyncIterable[reducer.Datum], md: reducer.Metadata + self, keys: list[str], datums: AsyncIterator[reducer.Datum], md: reducer.Metadata ) -> list[reducer.Message]: async for _ in datums: self.counter += 1 @@ -96,7 +96,7 @@ def test_reduce_server_creates_reducer_from_factory(tmp_path: Path): def test_reduce_server_propagates_handler_exception(tmp_path: Path): async def handler( - keys: list[str], datums: AsyncIterable[reducer.Datum], md: reducer.Metadata + keys: list[str], datums: AsyncIterator[reducer.Datum], md: reducer.Metadata ) -> list[reducer.Message]: async for _datum in datums: raise RuntimeError("reduce exploded") @@ -122,7 +122,7 @@ def factory(): def test_reduce_server_rejects_non_list_response(tmp_path: Path): - async def handler(keys: list[str], datums: AsyncIterable[reducer.Datum], md: reducer.Metadata) -> str: + async def handler(keys: list[str], datums: AsyncIterator[reducer.Datum], md: reducer.Metadata) -> str: async for _datum in datums: pass return "not messages" diff --git a/packages/pynumaflow-lite/tests/test_reducestream.py b/packages/pynumaflow-lite/tests/test_reducestream.py index 6d2ce0df..eecacb15 100644 --- a/packages/pynumaflow-lite/tests/test_reducestream.py +++ b/packages/pynumaflow-lite/tests/test_reducestream.py @@ -3,12 +3,14 @@ import pytest from _test_utils import run_python_server_with_rust_client +from pynumaflow_lite import reducestreamer SOCK_PATH = Path("/tmp/var/run/numaflow/reducestream.sock") SERVER_INFO = Path("/tmp/var/run/numaflow/reducestreamer-server-info") SCRIPTS = [ "reducestream_counter.py", + "reducestream_counter_func.py", ] @@ -20,3 +22,16 @@ def test_python_reducestream_server_and_rust_client(script: str, tmp_path: Path) server_info_path=SERVER_INFO, rust_bin_name="test_reducestream", ) + + +def test_reducestream_data_types_match_pythonic_api(): + datum = reducestreamer.Datum(keys=["key"], value=b"value") + assert datum.keys == ["key"] + assert datum.event_time + + message = reducestreamer.Message(b"value", keys=["key"]) + assert message == reducestreamer.Message(b"value", keys=["key"]) + assert reducestreamer.Message.to_drop().tags == ["U+005C__DROP__"] + + assert not hasattr(reducestreamer.Message, "message_to_drop") + assert not hasattr(reducestreamer, "PyAsyncDatumStream") diff --git a/packages/pynumaflow-lite/tests/test_reducestream_unit.py b/packages/pynumaflow-lite/tests/test_reducestream_unit.py new file mode 100644 index 00000000..e98a72c2 --- /dev/null +++ b/packages/pynumaflow-lite/tests/test_reducestream_unit.py @@ -0,0 +1,159 @@ +import asyncio +import functools +import os +import subprocess +import uuid +from collections.abc import AsyncIterator +from contextlib import suppress +from datetime import UTC, datetime +from pathlib import Path + +import pytest + +from pynumaflow_lite import reducestreamer + +CARGO_ROOT = Path(__file__).resolve().parent.parent + +WINDOW = reducestreamer.IntervalWindow(datetime(2023, 6, 1, tzinfo=UTC), datetime(2023, 6, 1, 0, 1, tzinfo=UTC)) + + +class Counter(reducestreamer.ReduceStreamer): + def __init__(self, initial: int = 0) -> None: + self.counter = initial + + async def handler( + self, keys: list[str], datums: AsyncIterator[reducestreamer.Datum], md: reducestreamer.Metadata + ) -> AsyncIterator[reducestreamer.Message]: + async for _ in datums: + self.counter += 1 + yield reducestreamer.Message(f"counter:{self.counter}".encode(), keys=keys) + + +async def _datums(count: int) -> AsyncIterator[reducestreamer.Datum]: + for _ in range(count): + yield reducestreamer.Datum(keys=["key"], value=b"value") + + +def test_reducestream_handler_can_be_unit_tested_without_server(): + md = reducestreamer.Metadata(interval_window=WINDOW) + + async def collect() -> list[reducestreamer.Message]: + return [message async for message in Counter(initial=10).handler(["key"], _datums(2), md)] + + messages = asyncio.run(collect()) + + assert messages == [ + reducestreamer.Message(b"counter:11", keys=["key"]), + reducestreamer.Message(b"counter:12", keys=["key"]), + ] + + +def test_reducestream_server_rejects_reducestreamer_instance(): + with pytest.raises(TypeError, match="not an instance"): + reducestreamer.ReduceStreamAsyncServer(Counter()) + + +async def _run_reducestream_client(sock_path: Path) -> subprocess.CompletedProcess[str]: + return await asyncio.to_thread( + subprocess.run, + ["cargo", "run", "--quiet", "--bin", "test_reducestream", "--", str(sock_path)], + cwd=CARGO_ROOT, + capture_output=True, + text=True, + timeout=60, + check=False, + ) + + +async def _exercise_server( + tmp_path: Path, + handler, + clients: list[subprocess.CompletedProcess[str]] | None = None, +) -> 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 = reducestreamer.ReduceStreamAsyncServer( + handler, + sock_file=str(sock_path), + server_info_file=str(server_info_path), + ) + + try: + async with server: + client = await _run_reducestream_client(sock_path) + if clients is not None: + clients.append(client) + finally: + for path in (sock_path, server_info_path): + with suppress(FileNotFoundError): + os.unlink(path) + + +def test_reducestream_server_creates_reducestreamer_from_factory(tmp_path: Path): + # The test client sends three datums and expects three streamed messages. + clients: list[subprocess.CompletedProcess[str]] = [] + asyncio.run(_exercise_server(tmp_path, functools.partial(Counter, initial=0), clients)) + + [client] = clients + assert client.returncode == 0, client.stderr + + +def test_reducestream_server_propagates_handler_exception(tmp_path: Path): + async def handler( + keys: list[str], datums: AsyncIterator[reducestreamer.Datum], md: reducestreamer.Metadata + ) -> AsyncIterator[reducestreamer.Message]: + async for _datum in datums: + raise RuntimeError("reducestream exploded") + yield reducestreamer.Message(b"unreachable") + + clients: list[subprocess.CompletedProcess[str]] = [] + with pytest.raises(RuntimeError, match="reducestream exploded"): + asyncio.run(_exercise_server(tmp_path, handler, clients)) + + # The client must get a gRPC error. Do not check for the handler text: + # numaflow-rs keeps only the first panic of the process. + [client] = clients + assert client.returncode != 0 + assert "UDF_EXECUTION_ERROR" in client.stderr + + +def test_reducestream_server_propagates_factory_exception(tmp_path: Path): + def factory(): + raise RuntimeError("factory exploded") + + with pytest.raises(RuntimeError, match="factory exploded"): + asyncio.run(_exercise_server(tmp_path, factory)) + + +def test_reducestream_server_rejects_non_message_item(tmp_path: Path): + async def handler( + keys: list[str], datums: AsyncIterator[reducestreamer.Datum], md: reducestreamer.Metadata + ) -> AsyncIterator[str]: + async for _datum in datums: + yield "not a message" + + with pytest.raises(TypeError, match="reducestream handler must yield Message, got str"): + asyncio.run(_exercise_server(tmp_path, handler)) + + +def test_reducestream_server_rejects_non_generator_handler(tmp_path: Path): + def factory(): + def handler(keys, datums, md): + return [] + + return handler + + with pytest.raises(TypeError, match="reducestream handler must be an async generator"): + asyncio.run(_exercise_server(tmp_path, factory)) + + +def test_reducestream_server_wait_ready_obeys_asyncio_timeout(tmp_path: Path): + server = reducestreamer.ReduceStreamAsyncServer(Counter, sock_file=str(tmp_path / "nobody-listens.sock")) + + async def wait() -> None: + async with asyncio.timeout(0.2): + await server.wait_ready() + + with pytest.raises(TimeoutError): + asyncio.run(wait())