pynumaflow-lite: Pythonic APIs for Sourcetransformer

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

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

packages/pynumaflow-lite/manifests/sourcetransform/Dockerfile

@@ -1,38 +1,43 @@
-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/
+FROM python:3.11-slim-trixie AS builder
 
-CMD ["python", "sourcetransform_event_filter.py"]
+COPY --from=ghcr.io/astral-sh/uv:0.12.13 /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
+
+# 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", "sourcetransform_event_filter.py"]

packages/pynumaflow-lite/manifests/sourcetransform/pyproject.toml

@@ -6,12 +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/sourcetransform/sourcetransform_event_filter.py

@@ -1,16 +1,13 @@
 import asyncio
-import signal
-from collections.abc import Callable
 from datetime import datetime, timezone
 
-from pynumaflow_lite import sourcetransformer
+from pynumaflow_lite.sourcetransformer import Datum, Message, SourceTransformAsyncServer, SourceTransformer
 
-# Define epoch timestamps for filtering
 january_first_2022 = datetime(2022, 1, 1, tzinfo=timezone.utc)
 january_first_2023 = datetime(2023, 1, 1, tzinfo=timezone.utc)
 
 
-class EventFilter(sourcetransformer.SourceTransformer):
+class EventFilter(SourceTransformer):
     """
     A source transformer that filters and routes messages based on event time.
 
@@ -19,65 +16,23 @@ class EventFilter(sourcetransformer.SourceTransformer):
     - Messages after 2022 are tagged with "after_year_2022"
     """
 
-    async def handler(self, keys: list[str], datum: sourcetransformer.Datum) -> sourcetransformer.Messages:
-        val = datum.value
+    async def handler(self, datum: Datum) -> list[Message]:
         event_time = datum.event_time
-        messages = sourcetransformer.Messages()
-
         if event_time < january_first_2022:
             print(f"Got event time: {event_time}, it is before 2022, so dropping")
-            messages.append(sourcetransformer.Message.message_to_drop(event_time))
-        elif event_time < january_first_2023:
-            print(f"Got event time: {event_time}, it is within year 2022, so forwarding to within_year_2022")
-            messages.append(
-                sourcetransformer.Message(
-                    value=val,
-                    event_time=january_first_2022,
-                    keys=keys,
-                    tags=["within_year_2022"],
-                )
-            )
-        else:
-            print(f"Got event time: {event_time}, it is after year 2022, so forwarding to after_year_2022")
-            messages.append(
-                sourcetransformer.Message(
-                    value=val,
-                    event_time=january_first_2023,
-                    keys=keys,
-                    tags=["after_year_2022"],
-                )
-            )
-
-        return messages
+            return [Message.to_drop(event_time)]
 
+        if event_time < january_first_2023:
+            print(f"Got event time: {event_time}, it is within year 2022, so forwarding to within_year_2022")
+            return [Message(datum.value, january_first_2022, keys=datum.keys, tags=["within_year_2022"])]
 
-async def start(
-    f: Callable[[list[str], sourcetransformer.Datum], sourcetransformer.Messages],
-):
-    server = sourcetransformer.SourceTransformAsyncServer()
+        print(f"Got event time: {event_time}, it is after year 2022, so forwarding to after_year_2022")
+        return [Message(datum.value, january_first_2023, keys=datum.keys, tags=["after_year_2022"])]
 
-    # 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:
+    await SourceTransformAsyncServer(EventFilter()).serve()
 
 
 if __name__ == "__main__":
-    async_handler = EventFilter()
-    asyncio.run(start(async_handler))
+    asyncio.run(main())

packages/pynumaflow-lite/pynumaflow_lite/__init__.py

@@ -18,6 +18,7 @@
 from ._sink_dtypes import Sinker
 from ._sink_server import SinkAsyncServer
 from ._source_dtypes import Sourcer
+from ._sourcetransform_server import SourceTransformAsyncServer
 from ._sourcetransformer_dtypes import SourceTransformer
 from .pynumaflow_lite import *  # noqa: F403  # Rust extension; exports resolved at runtime
 
@@ -51,6 +52,7 @@
 sinker.SinkAsyncServer = SinkAsyncServer
 sourcer.Sourcer = Sourcer
 sourcetransformer.SourceTransformer = SourceTransformer
+sourcetransformer.SourceTransformAsyncServer = SourceTransformAsyncServer
 sideinputer.SideInput = SideInput
 
 # Public API

packages/pynumaflow-lite/pynumaflow_lite/_sourcetransform_server.py

@@ -0,0 +1,125 @@
+from __future__ import annotations
+
+import asyncio
+import contextlib
+import signal
+from collections.abc import Awaitable, Callable
+from types import TracebackType
+from typing import TypeAlias
+
+from .pynumaflow_lite import sourcetransformer as _sourcetransformer
+
+Datum: TypeAlias = _sourcetransformer.Datum
+Message: TypeAlias = _sourcetransformer.Message
+
+_SHUTDOWN_SIGNALS = (signal.SIGINT, signal.SIGTERM)
+
+
+class SourceTransformAsyncServer:
+    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 = _sourcetransformer._SourceTransformAsyncServer(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 sourcetransform 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("sourcetransform 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("sourcetransform 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) -> SourceTransformAsyncServer:
+        """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("sourcetransform 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

packages/pynumaflow-lite/pynumaflow_lite/_sourcetransformer_dtypes.py

@@ -1,6 +1,6 @@
 from abc import ABCMeta, abstractmethod
 
-from pynumaflow_lite.sourcetransformer import Datum, Messages
+from pynumaflow_lite.sourcetransformer import Datum, Message
 
 
 class SourceTransformer(metaclass=ABCMeta):
@@ -20,16 +20,16 @@ class instance is sent as a callable.
         return self.handler(*args, **kwargs)
 
     @abstractmethod
-    async def handler(self, keys: list[str], datum: Datum) -> Messages:
+    async def handler(self, datum: Datum) -> list[Message]:
         """
         Implement this handler function which implements the SourceTransformer interface.
 
         Args:
-            keys: The keys associated with the message.
-            datum: The input datum containing value, event_time, watermark, and headers.
+            datum: The input datum containing keys, value, event_time, watermark,
+                headers, and metadata.
 
         Returns:
-            Messages: A collection of transformed messages with potentially modified
-                     event times and tags for conditional forwarding.
+            The transformed messages, with potentially modified event times and
+            tags for conditional forwarding.
         """
         pass