pynumaflow-lite: Pythonic APIs for Reducer and ReduceStreamer

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":

packages/pynumaflow-lite/manifests/reduce/.dockerignore

@@ -0,0 +1 @@
+.venv/

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"]

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"
-

packages/pynumaflow-lite/manifests/reduce/reduce_counter_class.py

@@ -1,49 +1,34 @@
 import asyncio
-import signal
-from collections.abc import AsyncIterable
+import functools
+from collections.abc import AsyncIterator
 
-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: AsyncIterator[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())

packages/pynumaflow-lite/manifests/reducestream/.dockerignore

@@ -0,0 +1 @@
+.venv/

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"]

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"
-

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())