pynumaflow-lite: Make mapper APIs more Pythonic.

packages/pynumaflow-lite/.gitignore

@@ -0,0 +1 @@
+.docker-cache/

packages/pynumaflow-lite/README.md

@@ -24,7 +24,12 @@ make test
 Go to `pynumaflow-lite` (top level) directory and run the below command.
 
 ```bash
-docker run --rm -v $(pwd):/io ghcr.io/pyo3/maturin build -i python3.11 --release 
+docker run --rm \
+  -v $(pwd):/io \
+  -v $(pwd)/.docker-cache/.cargo-docker-registry:/root/.cargo/registry \
+  -v $(pwd)/.docker-cache/cargo-docker-git:/root/.cargo/git \
+  -v $(pwd)/target:/io/target \
+  ghcr.io/pyo3/maturin build -i python3.11 --release
 ```
 
 This will create the `wheel` file in `target/wheels/` directory. You should copy it over to where we
@@ -33,5 +38,5 @@ are writing the python code referencing this library.
 e.g.,
 
 ```bash
-cp target/wheels/pynumaflow_lite-0.1.0-cp311-cp311-manylinux_2_17_aarch64.manylinux2014_aarch64.whl manifests/simple-async-map/
-```
+cp -v target/wheels/pynumaflow_lite-0.1.0a1-cp311-cp311-manylinux_2_17_aarch64.manylinux2014_aarch64.whl manifests/map/
+```

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

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

packages/pynumaflow-lite/manifests/map/Dockerfile

@@ -1,37 +1,42 @@
-FROM python:3.11-slim-bullseye 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 ./ ./
-
-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/
-
-CMD ["python", "map_cat.py"]
+FROM python:3.11-slim-trixie AS builder
+
+COPY --from=ghcr.io/astral-sh/uv:latest /uv /uvx /bin/
+
+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
+
+RUN uv pip install ./pynumaflow_lite-0.1.0a1-cp311-cp311-manylinux_2_17_aarch64.manylinux2014_aarch64.whl
+
+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", "map_cat.py"]

packages/pynumaflow-lite/manifests/map/README.md

@@ -16,4 +16,4 @@ k3d image import quay.io/numaio/numaflow/pynumaflow-lite-map-cat:v2
 
 ```bash
 kubectl apply -f pipeline.yaml
-```
+```

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

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

packages/pynumaflow-lite/manifests/map/pipeline.yaml

@@ -8,8 +8,8 @@ spec:
       source:
         # A self data generating source
         generator:
-          rpu: 500
-          duration: 1s
+          rpu: 1
+          duration: 3s
     - name: map
       partitions: 2
       scale:
@@ -27,4 +27,4 @@ spec:
     - from: in
       to: map
     - from: map
-      to: sink
+      to: sink

packages/pynumaflow-lite/manifests/map/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/pynumaflow_lite/__init__.py

@@ -66,6 +66,7 @@
 from ._accumulator_dtypes import Accumulator
 from ._batchmapper_dtypes import BatchMapper
 from ._map_dtypes import Mapper
+from ._map_server import MapAsyncServer
 from ._mapstream_dtypes import MapStreamer
 from ._reduce_dtypes import Reducer
 from ._reducestreamer_dtypes import ReduceStreamer
@@ -78,6 +79,7 @@
 
 if mapper is not None:
     mapper.Mapper = Mapper
+    mapper.MapAsyncServer = MapAsyncServer
 
 if batchmapper is not None:
     batchmapper.BatchMapper = BatchMapper

packages/pynumaflow-lite/pynumaflow_lite/_map_dtypes.py

@@ -1,6 +1,6 @@
 from abc import ABCMeta, abstractmethod
 
-from pynumaflow_lite.mapper import Datum, Messages
+from pynumaflow_lite.mapper import Datum, Message
 
 
 class Mapper(metaclass=ABCMeta):
@@ -17,7 +17,7 @@ class instance is sent as a callable.
         return self.handler(*args, **kwargs)
 
     @abstractmethod
-    async def handler(self, keys: list[str], payload: Datum) -> Messages:
+    async def handler(self, payload: Datum) -> list[Message]:
         """
         Implement this handler function which implements the MapAsyncCallable interface.
         """

packages/pynumaflow-lite/pynumaflow_lite/_map_server.py

@@ -0,0 +1,131 @@
+from __future__ import annotations
+
+import asyncio
+import contextlib
+import signal
+from collections.abc import Awaitable, Callable
+from types import TracebackType
+
+from .pynumaflow_lite import mapper as _mapper
+
+Datum = _mapper.Datum
+Message = _mapper.Message
+
+_SHUTDOWN_SIGNALS = (signal.SIGINT, signal.SIGTERM)
+
+
+class MapAsyncServer:
+    def __init__(
+        self,
+        handler: Callable[[Datum], Awaitable[list[Message]]],
+        *,
+        sock_file: str | None = None,
+        server_info_file: str | None = None,
+        install_signal_handlers: bool = True,
+    ) -> None:
+        self._core = _mapper._MapAsyncServer(sock_file, server_info_file)
+        self._handler = handler
+        self._install_signal_handlers = install_signal_handlers
+        self._task: asyncio.Task[None] | None = None
+        self._serving = False
+        self._installed_signals: list[signal.Signals] = []
+
+    async def serve(self) -> None:
+        """Run the map server until it stops.
+
+        This is the entrypoint for an application that already runs an event
+        loop. It returns when a shutdown signal arrives or when `stop()` runs.
+        """
+        await self._serve(install_signal_handlers=self._install_signal_handlers)
+
+    async def _serve(self, *, install_signal_handlers: bool) -> None:
+        if self._serving:
+            raise RuntimeError("map server is already serving")
+        self._serving = True
+        try:
+            if install_signal_handlers:
+                self._add_signal_handlers()
+            await self._core.start(self._handler)
+        finally:
+            self._remove_signal_handlers()
+            self._serving = False
+
+    def stop(self) -> None:
+        self._core.stop()
+
+    async def wait_ready(self, timeout: float = 30.0) -> None:
+        await self._core.wait_ready(timeout)
+
+    async def wait_for_termination(self) -> None:
+        """Wait until the background server task ends.
+
+        Use this inside an `async with` block. It raises the handler error if
+        the server task failed.
+        """
+        if self._task is None:
+            raise RuntimeError("map server is not serving")
+        await asyncio.shield(self._task)
+
+    def _add_signal_handlers(self) -> None:
+        loop = asyncio.get_running_loop()
+        for sig in _SHUTDOWN_SIGNALS:
+            try:
+                loop.add_signal_handler(sig, self.stop)
+            except (NotImplementedError, RuntimeError, OSError):
+                continue
+            self._installed_signals.append(sig)
+
+    def _remove_signal_handlers(self) -> None:
+        if not self._installed_signals:
+            return
+        try:
+            loop = asyncio.get_running_loop()
+        except RuntimeError:
+            self._installed_signals.clear()
+            return
+        for sig in self._installed_signals:
+            with contextlib.suppress(NotImplementedError, OSError):
+                loop.remove_signal_handler(sig)
+        self._installed_signals.clear()
+
+    async def __aenter__(self) -> MapAsyncServer:
+        """Start the server in a background task and wait until it is ready.
+
+        This form is for tests and for code that must run other work next to
+        the server. It never installs signal handlers.
+        """
+        if self._task is not None and not self._task.done():
+            raise RuntimeError("map server is already serving")
+
+        self._task = asyncio.create_task(self._serve(install_signal_handlers=False))
+        try:
+            await self.wait_ready()
+        except BaseException:
+            self.stop()
+            task, self._task = self._task, None
+            if task is not None:
+                # Surface the server error, if there is one. It explains the
+                # failure better than the `wait_ready` error does.
+                await task
+            raise
+        return self
+
+    async def __aexit__(
+        self,
+        exc_type: type[BaseException] | None,
+        exc: BaseException | None,
+        tb: TracebackType | None,
+    ) -> None:
+        self.stop()
+        if self._task is not None:
+            try:
+                await self._task
+            finally:
+                self._task = None
+
+    def run(self) -> None:
+        """Run the map server in a new event loop until it stops."""
+        try:
+            asyncio.run(self.serve())
+        except KeyboardInterrupt:
+            self.stop()