feat: sink integration for pynumaflow-lite (#281)

Signed-off-by: Vigith Maurice <[email protected]>

Author: vigith

packages/pynumaflow-lite/Cargo.toml

@@ -9,7 +9,7 @@ name = "pynumaflow_lite"
 crate-type = ["cdylib", "rlib"]
 
 [dependencies]
-numaflow = { git = "https://github.com/numaproj/numaflow-rs.git", branch = "export-accum-items" }
+numaflow = { git = "https://github.com/numaproj/numaflow-rs.git", rev = "fde3deafea634abbc347032ff409d33d4e1514b1" }
 pyo3 = { version = "0.26.0", features = ["chrono", "experimental-inspect"] }
 tokio = "1.47.1"
 tonic = "0.14.2"
@@ -22,6 +22,7 @@ pyo3-async-runtimes = { version = "0.26.0", features = ["tokio-runtime"] }
 futures-core = "0.3.31"
 pin-project = "1.1.10"
 
+## Binaries for testing
 
 [[bin]]
 name = "test_map"
@@ -31,12 +32,10 @@ path = "tests/bin/map.rs"
 name = "test_batchmap"
 path = "tests/bin/batchmap.rs"
 
-
 [[bin]]
 name = "test_mapstream"
 path = "tests/bin/mapstream.rs"
 
-
 [[bin]]
 name = "test_reduce"
 path = "tests/bin/reduce.rs"
@@ -48,3 +47,7 @@ path = "tests/bin/session_reduce.rs"
 [[bin]]
 name = "test_accumulator"
 path = "tests/bin/accumulator.rs"
+
+[[bin]]
+name = "test_sink"
+path = "tests/bin/sink.rs"

packages/pynumaflow-lite/manifests/sink/Dockerfile

@@ -0,0 +1,38 @@
+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", "sink_log.py"]
+

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

@@ -0,0 +1,28 @@
+To create the `wheel` file, refer [root](../../README.md)
+
+## HOWTO build Image
+
+```bash
+docker build . -t quay.io/numaio/numaflow/pynumaflow-lite-sink-log:v1 --load
+```
+
+### `k3d`
+
+Load it now to `k3d`
+
+```bash
+k3d image import quay.io/numaio/numaflow/pynumaflow-lite-sink-log:v1
+```
+
+### Minikube
+
+```bash
+minikube image load quay.io/numaio/numaflow/pynumaflow-lite-sink-log:v1
+```
+
+## Run the pipeline
+
+```bash
+kubectl apply -f pipeline.yaml
+```
+

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

@@ -0,0 +1,25 @@
+apiVersion: numaflow.numaproj.io/v1alpha1
+kind: Pipeline
+metadata:
+  name: simple-sink-log
+spec:
+  vertices:
+    - name: in
+      source:
+        # A self data generating source
+        generator:
+          rpu: 100
+          duration: 1s
+          msgSize: 8
+    - name: log-sink
+      scale:
+        min: 1
+      sink:
+        udsink:
+          container:
+            image: quay.io/numaio/numaflow/pynumaflow-lite-sink-log:v1
+            imagePullPolicy: Never
+  edges:
+    - from: in
+      to: log-sink
+

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

@@ -0,0 +1,17 @@
+[project]
+name = "sink-log"
+version = "0.1.0"
+description = "User-defined sink example using pynumaflow-lite"
+authors = [
+    { name = "Vigith Maurice", email = "[email protected]" }
+]
+readme = "README.md"
+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/sink/sink_log.py

@@ -0,0 +1,63 @@
+import asyncio
+import collections
+import logging
+import signal
+from collections.abc import AsyncIterable
+
+from pynumaflow_lite import sinker
+from pynumaflow_lite._sink_dtypes import Sinker
+
+# Configure logging
+logging.basicConfig(level=logging.INFO)
+_LOGGER = logging.getLogger(__name__)
+
+
+class SimpleLogSink(Sinker):
+    """
+    Simple log sink that logs each message and returns success responses.
+    """
+
+    async def handler(self, datums: AsyncIterable[sinker.Datum]) -> sinker.Responses:
+        responses = sinker.Responses()
+        async for msg in datums:
+            _LOGGER.info("User Defined Sink: %s", msg.value.decode("utf-8"))
+            responses.append(sinker.Response.as_success(msg.id))
+            # if we are not able to write to sink and if we have a fallback sink configured
+            # we can use Response.as_fallback(msg.id) to write the message to fallback sink
+        return responses
+
+
+# Optional: ensure default signal handlers are in place so asyncio.run can handle them cleanly.
+signal.signal(signal.SIGINT, signal.default_int_handler)
+try:
+    signal.signal(signal.SIGTERM, signal.SIG_DFL)
+except AttributeError:
+    pass
+
+
+async def start(f: collections.abc.Callable):
+    server = sinker.SinkAsyncServer()
+
+    # Register loop-level signal handlers so we control shutdown and avoid asyncio.run
+    loop = asyncio.get_running_loop()
+    try:
+        loop.add_signal_handler(signal.SIGINT, lambda: server.stop())
+        loop.add_signal_handler(signal.SIGTERM, lambda: server.stop())
+    except (NotImplementedError, RuntimeError):
+        pass
+
+    try:
+        await server.start(f)
+        print("Shutting down gracefully...")
+    except asyncio.CancelledError:
+        try:
+            server.stop()
+        except Exception:
+            pass
+        return
+
+
+if __name__ == "__main__":
+    async_handler = SimpleLogSink()
+    asyncio.run(start(async_handler))
+

packages/pynumaflow-lite/pynumaflow_lite/__init__.pyi

@@ -8,5 +8,6 @@ from . import mapstreamer as mapstreamer
 from . import reducer as reducer
 from . import session_reducer as session_reducer
 from . import accumulator as accumulator
+from . import sinker as sinker
 
-__all__ = ['mapper', 'batchmapper', 'mapstreamer', 'reducer', 'session_reducer', 'accumulator']
+__all__ = ['mapper', 'batchmapper', 'mapstreamer', 'reducer', 'session_reducer', 'accumulator', 'sinker']

packages/pynumaflow-lite/pynumaflow_lite/_sink_dtypes.py

@@ -0,0 +1,21 @@
+from abc import ABCMeta, abstractmethod
+from pynumaflow_lite.sinker import Datum, Responses
+from collections.abc import AsyncIterable
+
+
+class Sinker(metaclass=ABCMeta):
+    """
+    Provides an interface to write a Sink servicer.
+    """
+
+    def __call__(self, *args, **kwargs):
+        return self.handler(*args, **kwargs)
+
+    @abstractmethod
+    async def handler(self, datums: AsyncIterable[Datum]) -> Responses:
+        """
+        Implement this handler function for sink.
+        Process the stream of datums and return responses.
+        """
+        pass
+

packages/pynumaflow-lite/pynumaflow_lite/sinker.pyi

@@ -0,0 +1,92 @@
+from __future__ import annotations
+
+from typing import Optional, List, Dict, Callable, Awaitable, Any, AsyncIterator
+import datetime as _dt
+
+
+class KeyValueGroup:
+    key_value: Dict[str, bytes]
+
+    def __init__(self, key_value: Optional[Dict[str, bytes]] = ...) -> None: ...
+
+    @staticmethod
+    def from_dict(key_value: Dict[str, bytes]) -> KeyValueGroup: ...
+
+
+class Message:
+    keys: Optional[List[str]]
+    value: bytes
+    user_metadata: Optional[Dict[str, KeyValueGroup]]
+
+    def __init__(
+            self,
+            value: bytes,
+            keys: Optional[List[str]] = ...,
+            user_metadata: Optional[Dict[str, KeyValueGroup]] = ...,
+    ) -> None: ...
+
+
+class Response:
+    id: str
+
+    @staticmethod
+    def as_success(id: str) -> Response: ...
+
+    @staticmethod
+    def as_failure(id: str, err_msg: str) -> Response: ...
+
+    @staticmethod
+    def as_fallback(id: str) -> Response: ...
+
+    @staticmethod
+    def as_serve(id: str, payload: bytes) -> Response: ...
+
+    @staticmethod
+    def as_on_success(id: str, message: Optional[Message] = ...) -> Response: ...
+
+
+class Responses:
+    def __init__(self) -> None: ...
+
+    def append(self, response: Response) -> None: ...
+
+
+class Datum:
+    keys: List[str]
+    value: bytes
+    watermark: _dt.datetime
+    eventtime: _dt.datetime
+    id: str
+    headers: Dict[str, str]
+
+    def __repr__(self) -> str: ...
+
+    def __str__(self) -> str: ...
+
+
+class SinkAsyncServer:
+    def __init__(
+            self,
+            sock_file: str | None = ...,
+            info_file: str | None = ...,
+    ) -> None: ...
+
+    def start(self, py_func: Callable[..., Any]) -> Awaitable[None]: ...
+
+    def stop(self) -> None: ...
+
+
+class Sinker:
+    async def handler(self, datums: AsyncIterator[Datum]) -> Responses: ...
+
+
+__all__ = [
+    "KeyValueGroup",
+    "Message",
+    "Response",
+    "Responses",
+    "Datum",
+    "SinkAsyncServer",
+    "Sinker",
+]
+

packages/pynumaflow-lite/src/lib.rs

@@ -6,6 +6,7 @@ pub mod pyiterables;
 pub mod pyrs;
 pub mod reduce;
 pub mod session_reduce;
+pub mod sink;
 
 use pyo3::prelude::*;
 
@@ -51,6 +52,13 @@ fn accumulator(_py: Python, m: &Bound<PyModule>) -> PyResult<()> {
     Ok(())
 }
 
+/// Submodule: pynumaflow_lite.sinker
+#[pymodule]
+fn sinker(_py: Python, m: &Bound<PyModule>) -> PyResult<()> {
+    crate::sink::populate_py_module(m)?;
+    Ok(())
+}
+
 /// Top-level Python module `pynumaflow_lite` with submodules like `mapper`, `batchmapper`, and `mapstreamer`.
 #[pymodule]
 fn pynumaflow_lite(py: Python, m: &Bound<PyModule>) -> PyResult<()> {
@@ -61,6 +69,7 @@ fn pynumaflow_lite(py: Python, m: &Bound<PyModule>) -> PyResult<()> {
     m.add_wrapped(pyo3::wrap_pymodule!(reducer))?;
     m.add_wrapped(pyo3::wrap_pymodule!(session_reducer))?;
     m.add_wrapped(pyo3::wrap_pymodule!(accumulator))?;
+    m.add_wrapped(pyo3::wrap_pymodule!(sinker))?;
 
     // Ensure it's importable as `pynumaflow_lite.mapper` as well as attribute access
     let binding = m.getattr("mapper")?;
@@ -116,5 +125,14 @@ fn pynumaflow_lite(py: Python, m: &Bound<PyModule>) -> PyResult<()> {
         .getattr("modules")?
         .set_item(fullname, &sub)?;
 
+    // Ensure it's importable as `pynumaflow_lite.sinker` as well
+    let binding = m.getattr("sinker")?;
+    let sub = binding.downcast::<PyModule>()?;
+    let fullname = "pynumaflow_lite.sinker";
+    sub.setattr("__name__", fullname)?;
+    py.import("sys")?
+        .getattr("modules")?
+        .set_item(fullname, &sub)?;
+
     Ok(())
 }