From 028459c8da13b92291a8a66469edfcabd61a7f45 Mon Sep 17 00:00:00 2001 From: Vigith Maurice Date: Sat, 8 Nov 2025 09:26:11 -0800 Subject: [PATCH 1/7] chore: update to latest rust-sdk commit Signed-off-by: Vigith Maurice --- packages/pynumaflow-lite/Cargo.toml | 6 +++++- 1 file changed, 5 insertions(+), 1 deletion(-) diff --git a/packages/pynumaflow-lite/Cargo.toml b/packages/pynumaflow-lite/Cargo.toml index 032f8b68..b7414edb 100644 --- a/packages/pynumaflow-lite/Cargo.toml +++ b/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" @@ -48,3 +48,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" From 500995d1ca5537fbcf2844c9c904541443515427 Mon Sep 17 00:00:00 2001 From: Vigith Maurice Date: Sat, 8 Nov 2025 09:27:06 -0800 Subject: [PATCH 2/7] feat: sink ffi integration Signed-off-by: Vigith Maurice --- packages/pynumaflow-lite/src/lib.rs | 18 + packages/pynumaflow-lite/src/sink/mod.rs | 401 ++++++++++++++++++++ packages/pynumaflow-lite/src/sink/server.rs | 106 ++++++ 3 files changed, 525 insertions(+) create mode 100644 packages/pynumaflow-lite/src/sink/mod.rs create mode 100644 packages/pynumaflow-lite/src/sink/server.rs diff --git a/packages/pynumaflow-lite/src/lib.rs b/packages/pynumaflow-lite/src/lib.rs index 89339eb5..c5f7af21 100644 --- a/packages/pynumaflow-lite/src/lib.rs +++ b/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) -> PyResult<()> { Ok(()) } +/// Submodule: pynumaflow_lite.sinker +#[pymodule] +fn sinker(_py: Python, m: &Bound) -> 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) -> PyResult<()> { @@ -61,6 +69,7 @@ fn pynumaflow_lite(py: Python, m: &Bound) -> 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) -> 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::()?; + let fullname = "pynumaflow_lite.sinker"; + sub.setattr("__name__", fullname)?; + py.import("sys")? + .getattr("modules")? + .set_item(fullname, &sub)?; + Ok(()) } diff --git a/packages/pynumaflow-lite/src/sink/mod.rs b/packages/pynumaflow-lite/src/sink/mod.rs new file mode 100644 index 00000000..6ccac8e7 --- /dev/null +++ b/packages/pynumaflow-lite/src/sink/mod.rs @@ -0,0 +1,401 @@ +use std::collections::HashMap; + +use numaflow::sink; + +use chrono::{DateTime, Utc}; + +/// Sink interface managed by Python. Python code will start the server +/// and can pass in the Python coroutine. +pub mod server; + +use tokio::sync::mpsc; + +use pyo3::prelude::*; +use std::sync::Mutex; + +/// KeyValueGroup represents a group of key-value pairs for user metadata. +#[pyclass(module = "pynumaflow_lite.sinker")] +#[derive(Clone, Default, Debug)] +pub struct KeyValueGroup { + pub key_value: HashMap>, +} + +#[pymethods] +impl KeyValueGroup { + #[new] + #[pyo3(signature = (key_value: "dict[str, bytes] | None"=None) -> "KeyValueGroup")] + fn new(key_value: Option>>) -> Self { + Self { + key_value: key_value.unwrap_or_default(), + } + } + + /// Create a KeyValueGroup from a dictionary of string to bytes. + #[staticmethod] + #[pyo3(signature = (key_value: "dict[str, bytes]") -> "KeyValueGroup")] + fn from_dict(key_value: HashMap>) -> Self { + Self { key_value } + } +} + +impl From for sink::KeyValueGroup { + fn from(value: KeyValueGroup) -> Self { + Self { + key_value: value.key_value, + } + } +} + +/// Message for OnSuccess sink response. +/// Contains information that needs to be sent to the OnSuccess sink. +#[pyclass(module = "pynumaflow_lite.sinker")] +#[derive(Clone, Default, Debug)] +pub struct Message { + pub keys: Option>, + pub value: Vec, + pub user_metadata: Option>, +} + +#[pymethods] +impl Message { + /// Create a new Message with the given value. + /// Keys and user_metadata are optional. + #[new] + #[pyo3(signature = (value: "bytes", keys: "list[str] | None"=None, user_metadata: "dict[str, KeyValueGroup] | None"=None) -> "Message")] + fn new( + value: Vec, + keys: Option>, + user_metadata: Option>, + ) -> Self { + Self { + value, + keys, + user_metadata, + } + } +} + +impl From for sink::Message { + fn from(value: Message) -> Self { + Self { + keys: value.keys, + value: value.value, + user_metadata: value + .user_metadata + .map(|m| m.into_iter().map(|(k, v)| (k, v.into())).collect()), + } + } +} + +/// Response for a single datum in the sink. +#[pyclass(module = "pynumaflow_lite.sinker")] +#[derive(Clone, Debug)] +pub struct Response { + pub id: String, + pub response_type: ResponseType, + pub err: Option, + pub serve_response: Option>, + pub on_success_msg: Option, +} + +#[pymethods] +impl Response { + /// Create a success response. + #[staticmethod] + #[pyo3(signature = (id: "str") -> "Response")] + fn as_success(id: String) -> Self { + Self { + id, + response_type: ResponseType::Success, + err: None, + serve_response: None, + on_success_msg: None, + } + } + + /// Create a failure response with an error message. + #[staticmethod] + #[pyo3(signature = (id: "str", err_msg: "str") -> "Response")] + fn as_failure(id: String, err_msg: String) -> Self { + Self { + id, + response_type: ResponseType::Failure, + err: Some(err_msg), + serve_response: None, + on_success_msg: None, + } + } + + /// Create a fallback response to forward to fallback sink. + #[staticmethod] + #[pyo3(signature = (id: "str") -> "Response")] + fn as_fallback(id: String) -> Self { + Self { + id, + response_type: ResponseType::Fallback, + err: None, + serve_response: None, + on_success_msg: None, + } + } + + /// Create a serve response with payload for serving store. + #[staticmethod] + #[pyo3(signature = (id: "str", payload: "bytes") -> "Response")] + fn as_serve(id: String, payload: Vec) -> Self { + Self { + id, + response_type: ResponseType::Serve, + err: None, + serve_response: Some(payload), + on_success_msg: None, + } + } + + /// Create an OnSuccess response with optional message. + /// If message is None, the original message will be sent to onSuccess sink. + #[staticmethod] + #[pyo3(signature = (id: "str", message: "Message | None"=None) -> "Response")] + fn as_on_success(id: String, message: Option) -> Self { + Self { + id, + response_type: ResponseType::OnSuccess, + err: None, + serve_response: None, + on_success_msg: message, + } + } +} + +/// Internal enum to track response type +#[derive(Clone, Debug)] +pub enum ResponseType { + Success, + Failure, + Fallback, + Serve, + OnSuccess, +} + +impl From for sink::Response { + fn from(value: Response) -> Self { + let response_type = match value.response_type { + ResponseType::Success => sink::ResponseType::Success, + ResponseType::Failure => sink::ResponseType::Failure, + ResponseType::Fallback => sink::ResponseType::FallBack, + ResponseType::Serve => sink::ResponseType::Serve, + ResponseType::OnSuccess => sink::ResponseType::OnSuccess, + }; + + Self { + id: value.id, + response_type, + err: value.err, + serve_response: value.serve_response, + on_success_msg: value.on_success_msg.map(|m| m.into()), + } + } +} + +/// A collection of Response objects. +#[pyclass(module = "pynumaflow_lite.sinker")] +#[derive(Clone, Debug)] +pub struct Responses { + pub(crate) responses: Vec, +} + +#[pymethods] +impl Responses { + #[new] + #[pyo3(signature = () -> "Responses")] + fn new() -> Self { + Self { responses: vec![] } + } + + /// Append a Response to the collection. + #[pyo3(signature = (response: "Response"))] + fn append(&mut self, response: Response) { + self.responses.push(response); + } + + fn __repr__(&self) -> String { + format!("Responses(count={})", self.responses.len()) + } + + fn __str__(&self) -> String { + format!("Responses(count={})", self.responses.len()) + } +} + +/// The incoming Datum for Sink +#[pyclass(module = "pynumaflow_lite.sinker")] +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 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, + /// ID is the unique id of the message to be sent to the Sink. + #[pyo3(get)] + pub id: String, + /// Headers for the message. + #[pyo3(get)] + pub headers: HashMap, +} + +impl Datum { + fn new( + keys: Vec, + value: Vec, + watermark: DateTime, + eventtime: DateTime, + id: String, + headers: HashMap, + ) -> Self { + Self { + keys, + value, + watermark, + eventtime, + id, + headers, + } + } + + fn __repr__(&self) -> String { + format!( + "Datum(keys={:?}, value={:?}, watermark={}, eventtime={}, id={}, headers={:?})", + self.keys, self.value, self.watermark, self.eventtime, self.id, self.headers + ) + } + + fn __str__(&self) -> String { + format!( + "Datum(keys={:?}, value={:?}, watermark={}, eventtime={}, id={}, headers={:?})", + self.keys, + String::from_utf8_lossy(&self.value), + self.watermark, + self.eventtime, + self.id, + self.headers + ) + } +} + +impl From for Datum { + fn from(value: sink::SinkRequest) -> Self { + Datum::new( + value.keys, + value.value, + value.watermark, + value.event_time, + value.id, + value.headers, + ) + } +} + +/// Python-visible async iterator that yields Datum items from a Tokio mpsc channel. +/// This is a thin wrapper around the generic AsyncChannelStream implementation. +#[pyclass(module = "pynumaflow_lite.sinker")] +pub struct PyAsyncDatumStream { + inner: crate::pyiterables::AsyncChannelStream, +} + +#[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 + } + + fn __anext__<'a>(&self, py: Python<'a>) -> PyResult> { + self.inner.py_anext(py) + } +} + +impl PyAsyncDatumStream { + pub fn new_with(rx: mpsc::Receiver) -> Self { + Self { + inner: crate::pyiterables::AsyncChannelStream::new(rx), + } + } +} + +/// Async Sink Server that can be started from Python code +#[pyclass(module = "pynumaflow_lite.sinker")] +pub struct SinkAsyncServer { + sock_file: String, + info_file: String, + shutdown_tx: Mutex>>, +} + +#[pymethods] +impl SinkAsyncServer { + #[new] + #[pyo3(signature = (sock_file: "str | None"=sink::SOCK_ADDR.to_string(), info_file: "str | None"=sink::SERVER_INFO_FILE.to_string()) -> "SinkAsyncServer" + )] + fn new(sock_file: String, info_file: String) -> Self { + Self { + sock_file, + info_file, + shutdown_tx: Mutex::new(None), + } + } + + /// Start the server with the given Python function. + #[pyo3(signature = (py_func: "callable") -> "None")] + pub fn start<'a>(&self, py: Python<'a>, py_func: Py) -> PyResult> { + let sock_file = self.sock_file.clone(); + let info_file = self.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, async move { + crate::sink::server::start(py_func, sock_file, info_file, rx) + .await + .expect("server failed to start"); + 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(()); + } + Ok(()) + } +} + +/// Helper to populate a PyModule with sink 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/sink/server.rs b/packages/pynumaflow-lite/src/sink/server.rs new file mode 100644 index 00000000..937eb9f8 --- /dev/null +++ b/packages/pynumaflow-lite/src/sink/server.rs @@ -0,0 +1,106 @@ +use numaflow::sink; +use numaflow::shared::ServerExtras; + +use pyo3::prelude::*; +use std::sync::Arc; + +pub(crate) struct PySinkRunner { + pub(crate) event_loop: Arc>, + pub(crate) py_func: Arc>, +} + +#[tonic::async_trait] +impl sink::Sinker for PySinkRunner { + async fn sink( + &self, + mut input: tokio::sync::mpsc::Receiver, + ) -> Vec { + // Create a channel to stream Datum into Python as an async iterator + let (tx, rx) = tokio::sync::mpsc::channel::(64); + + // Spawn a task forwarding incoming datums to the Python-facing channel + let forwarder = tokio::spawn(async move { + while let Some(req) = input.recv().await { + if tx.send(req.into()).await.is_err() { + break; + } + } + // When input ends, dropping tx closes the channel + }); + + // Call the Python coroutine: py_func(datums: AsyncIterable[Datum]) -> Responses + let fut = Python::attach(|py| { + let locals = pyo3_async_runtimes::TaskLocals::new(self.event_loop.bind(py).clone()); + let py_func = self.py_func.clone(); + + let stream = crate::sink::PyAsyncDatumStream::new_with(rx); + let coro = py_func.call1(py, (stream,)).unwrap().into_bound(py); + pyo3_async_runtimes::into_future_with_locals(&locals, coro).unwrap() + }); + + let result = fut.await.unwrap(); + + // Ensure forwarder completes + let _ = forwarder.await; + + let responses = Python::attach(|py| { + let x: crate::sink::Responses = result.extract(py).unwrap(); + x + }); + + responses + .responses + .into_iter() + .map(|resp| resp.into()) + .collect::>() + } +} + +/// Start the sink server by spinning up a dedicated Python asyncio loop and wiring shutdown. +pub(super) async fn start( + py_func: 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 event_loop = rx.await.unwrap(); + + let (sig_handle, combined_rx) = crate::pyrs::setup_sig_handler(shutdown_rx); + + let py_runner = PySinkRunner { + py_func: Arc::new(py_func), + event_loop: event_loop.clone(), + }; + + let server = numaflow::sink::Server::new(py_runner) + .with_socket_file(sock_file) + .with_server_info_file(info_file); + + let result = server + .start_with_shutdown(combined_rx) + .await + .map_err(|e| pyo3::PyErr::new::(e.to_string())); + + // Ensure the event loop is stopped even if shutdown came from elsewhere. + Python::attach(|py| { + if let Ok(stop_cb) = event_loop.getattr(py, "stop") { + let _ = event_loop.call_method1(py, "call_soon_threadsafe", (stop_cb,)); + } + }); + + println!("Numaflow Sink 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"); + let _ = sig_handle.abort(); + } + + result +} + From 8c411e7fdac1ea4342f5c9c3d8ddc68d64e55b66 Mon Sep 17 00:00:00 2001 From: Vigith Maurice Date: Sat, 8 Nov 2025 09:27:28 -0800 Subject: [PATCH 3/7] feat: typing hints Signed-off-by: Vigith Maurice --- .../pynumaflow_lite/__init__.pyi | 3 +- .../pynumaflow_lite/_sink_dtypes.py | 21 +++++ .../pynumaflow_lite/sinker.pyi | 92 +++++++++++++++++++ 3 files changed, 115 insertions(+), 1 deletion(-) create mode 100644 packages/pynumaflow-lite/pynumaflow_lite/_sink_dtypes.py create mode 100644 packages/pynumaflow-lite/pynumaflow_lite/sinker.pyi diff --git a/packages/pynumaflow-lite/pynumaflow_lite/__init__.pyi b/packages/pynumaflow-lite/pynumaflow_lite/__init__.pyi index 7bcf17db..d7e03a7b 100644 --- a/packages/pynumaflow-lite/pynumaflow_lite/__init__.pyi +++ b/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'] diff --git a/packages/pynumaflow-lite/pynumaflow_lite/_sink_dtypes.py b/packages/pynumaflow-lite/pynumaflow_lite/_sink_dtypes.py new file mode 100644 index 00000000..479ad60b --- /dev/null +++ b/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 + diff --git a/packages/pynumaflow-lite/pynumaflow_lite/sinker.pyi b/packages/pynumaflow-lite/pynumaflow_lite/sinker.pyi new file mode 100644 index 00000000..1bb5b9d9 --- /dev/null +++ b/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", +] + From 0fac9a215aa68ca7c1071d3bafa7306201a61d35 Mon Sep 17 00:00:00 2001 From: Vigith Maurice Date: Sat, 8 Nov 2025 09:27:47 -0800 Subject: [PATCH 4/7] test: tests for sink Signed-off-by: Vigith Maurice --- packages/pynumaflow-lite/tests/bin/sink.rs | 120 ++++++++++++++++++ .../tests/examples/sink_log.py | 52 ++++++++ .../tests/examples/sink_log_class.py | 59 +++++++++ packages/pynumaflow-lite/tests/test_sink.py | 24 ++++ 4 files changed, 255 insertions(+) create mode 100644 packages/pynumaflow-lite/tests/bin/sink.rs create mode 100644 packages/pynumaflow-lite/tests/examples/sink_log.py create mode 100644 packages/pynumaflow-lite/tests/examples/sink_log_class.py create mode 100644 packages/pynumaflow-lite/tests/test_sink.py diff --git a/packages/pynumaflow-lite/tests/bin/sink.rs b/packages/pynumaflow-lite/tests/bin/sink.rs new file mode 100644 index 00000000..4d1342ce --- /dev/null +++ b/packages/pynumaflow-lite/tests/bin/sink.rs @@ -0,0 +1,120 @@ +use std::collections::HashMap; +use std::env; +use std::path::PathBuf; + +use numaflow::proto; +use numaflow::proto::sink::sink_client::SinkClient; +use tokio::net::UnixStream; +use tokio::sync::mpsc; +use tokio_stream::wrappers::ReceiverStream; +use tonic::transport::Uri; +use tower::service_fn; + +#[tokio::main] +async fn main() -> Result<(), Box> { + // Allow overriding the socket path via first CLI arg or env var. + let sock_path = env::args() + .nth(1) + .or_else(|| env::var("NUMAFLOW_SINK_SOCK").ok()) + .unwrap_or_else(|| "/tmp/var/run/numaflow/sink.sock".to_string()); + + // Set up tonic channel over Unix Domain Socket. + let channel = tonic::transport::Endpoint::try_from("http://[::]:50051")? + .connect_with_connector(service_fn(move |_: Uri| { + let sock = PathBuf::from(sock_path.clone()); + async move { + Ok::<_, std::io::Error>(hyper_util::rt::TokioIo::new( + UnixStream::connect(sock).await?, + )) + } + })) + .await?; + + let mut client = SinkClient::new(channel); + + let (tx, rx) = mpsc::channel(16); + + // Handshake to initialize the stream + let handshake_request = proto::sink::SinkRequest { + request: None, + handshake: Some(proto::sink::Handshake { sot: true }), + status: None, + }; + tx.send(handshake_request).await.unwrap(); + + let resp = client.sink_fn(ReceiverStream::new(rx)).await.unwrap(); + let mut resp = resp.into_inner(); + + // Expect handshake response from server + let handshake_response = resp.message().await.unwrap(); + assert!(handshake_response.is_some()); + let handshake_response = handshake_response.unwrap(); + assert!(handshake_response.handshake.is_some()); + + // Build three requests with IDs + let mk_req = |id: &str, keys: Vec<&str>, value: &str| -> proto::sink::SinkRequest { + proto::sink::SinkRequest { + request: Some(proto::sink::sink_request::Request { + keys: keys.into_iter().map(|s| s.to_string()).collect(), + value: value.as_bytes().to_vec(), + watermark: Some(prost_types::Timestamp::default()), + event_time: Some(prost_types::Timestamp::default()), + id: id.to_string(), + headers: HashMap::new(), + }), + handshake: None, + status: None, + } + }; + + let req1 = mk_req("id-1", vec!["k1"], "hello-1"); + let req2 = mk_req("id-2", vec!["k2"], "hello-2"); + let req3 = mk_req("id-3", vec!["k3"], "hello-3"); + + // Sender must live long enough; keep a clone for clarity + let tx_ref = tx; + tx_ref.send(req1).await.unwrap(); + tx_ref.send(req2).await.unwrap(); + tx_ref.send(req3).await.unwrap(); + + // Send End-Of-Transmission marker via status + let eot = numaflow::proto::sink::TransmissionStatus { eot: true }; + let eot_req = proto::sink::SinkRequest { + request: None, + handshake: None, + status: Some(eot), + }; + tx_ref.send(eot_req).await.unwrap(); + + // Collect the batch response + let maybe = resp.message().await.unwrap(); + assert!(maybe.is_some()); + let r = maybe.unwrap(); + + // Verify we got 3 responses + assert_eq!(r.results.len(), 3); + + // Verify all responses are successful + for result in &r.results { + assert_eq!(result.status, proto::sink::Status::Success as i32); + assert_eq!(result.err_msg, ""); + } + + // Verify the IDs match + let ids: Vec = r.results.iter().map(|res| res.id.clone()).collect(); + assert!(ids.contains(&"id-1".to_string())); + assert!(ids.contains(&"id-2".to_string())); + assert!(ids.contains(&"id-3".to_string())); + + // Expect EOT response + let eot_response = resp.message().await.unwrap(); + assert!(eot_response.is_some()); + let eot_response = eot_response.unwrap(); + assert!(eot_response.status.is_some()); + assert!(eot_response.status.unwrap().eot); + + println!("All sink tests passed!"); + + Ok(()) +} + diff --git a/packages/pynumaflow-lite/tests/examples/sink_log.py b/packages/pynumaflow-lite/tests/examples/sink_log.py new file mode 100644 index 00000000..6d11382a --- /dev/null +++ b/packages/pynumaflow-lite/tests/examples/sink_log.py @@ -0,0 +1,52 @@ +import asyncio +import collections.abc +import logging +import signal + +from pynumaflow_lite import sinker + +# Configure logging +logging.basicConfig(level=logging.INFO) +_LOGGER = logging.getLogger(__name__) + + +async def async_handler(datums: collections.abc.AsyncIterator[sinker.Datum]) -> sinker.Responses: + """ + Simple log sink that logs each message and returns success 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 + + +async def start(f: callable): + sock_file = "/tmp/var/run/numaflow/sink.sock" + server_info_file = "/tmp/var/run/numaflow/sinker-server-info" + server = sinker.SinkAsyncServer(sock_file, server_info_file) + + # Register loop-level signal handlers to request graceful shutdown + 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__": + asyncio.run(start(async_handler)) + diff --git a/packages/pynumaflow-lite/tests/examples/sink_log_class.py b/packages/pynumaflow-lite/tests/examples/sink_log_class.py new file mode 100644 index 00000000..6df2c4d6 --- /dev/null +++ b/packages/pynumaflow-lite/tests/examples/sink_log_class.py @@ -0,0 +1,59 @@ +import asyncio +import logging +import signal + +from pynumaflow_lite import sinker +from pynumaflow_lite._sink_dtypes import Sinker +from collections.abc import AsyncIterable + +# 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. + This is the class-based approach matching the user's example. + """ + + 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 + + +async def start(): + sock_file = "/tmp/var/run/numaflow/sink.sock" + server_info_file = "/tmp/var/run/numaflow/sinker-server-info" + server = sinker.SinkAsyncServer(sock_file, server_info_file) + + # Create an instance of the sink handler + handler = SimpleLogSink() + + # Register loop-level signal handlers to request graceful shutdown + 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: + try: + server.stop() + except Exception: + pass + return + + +if __name__ == "__main__": + asyncio.run(start()) + diff --git a/packages/pynumaflow-lite/tests/test_sink.py b/packages/pynumaflow-lite/tests/test_sink.py new file mode 100644 index 00000000..3890e416 --- /dev/null +++ b/packages/pynumaflow-lite/tests/test_sink.py @@ -0,0 +1,24 @@ +from pathlib import Path + +import pytest + +from _test_utils import run_python_server_with_rust_client + +SOCK_PATH = Path("/tmp/var/run/numaflow/sink.sock") +SERVER_INFO = Path("/tmp/var/run/numaflow/sinker-server-info") + +SCRIPTS = [ + "sink_log.py", + "sink_log_class.py", +] + + +@pytest.mark.parametrize("script", SCRIPTS) +def test_python_sink_server_and_rust_client(script: str, tmp_path: Path): + run_python_server_with_rust_client( + script=script, + sock_path=SOCK_PATH, + server_info_path=SERVER_INFO, + rust_bin_name="test_sink", + ) + From fc1d5b32e9587275232e445e523126f86a2bb1b3 Mon Sep 17 00:00:00 2001 From: Vigith Maurice Date: Sat, 8 Nov 2025 09:51:55 -0800 Subject: [PATCH 5/7] feat: add manifests Signed-off-by: Vigith Maurice --- packages/pynumaflow-lite/Cargo.toml | 3 +- .../pynumaflow-lite/manifests/sink/Dockerfile | 38 ++++++++++++ .../pynumaflow-lite/manifests/sink/README.md | 46 ++++++++++++++ .../manifests/sink/pipeline.yaml | 25 ++++++++ .../manifests/sink/pyproject.toml | 17 +++++ .../manifests/sink/sink_log.py | 62 +++++++++++++++++++ 6 files changed, 189 insertions(+), 2 deletions(-) create mode 100644 packages/pynumaflow-lite/manifests/sink/Dockerfile create mode 100644 packages/pynumaflow-lite/manifests/sink/README.md create mode 100644 packages/pynumaflow-lite/manifests/sink/pipeline.yaml create mode 100644 packages/pynumaflow-lite/manifests/sink/pyproject.toml create mode 100644 packages/pynumaflow-lite/manifests/sink/sink_log.py diff --git a/packages/pynumaflow-lite/Cargo.toml b/packages/pynumaflow-lite/Cargo.toml index b7414edb..5182f0e6 100644 --- a/packages/pynumaflow-lite/Cargo.toml +++ b/packages/pynumaflow-lite/Cargo.toml @@ -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" diff --git a/packages/pynumaflow-lite/manifests/sink/Dockerfile b/packages/pynumaflow-lite/manifests/sink/Dockerfile new file mode 100644 index 00000000..90ba07c7 --- /dev/null +++ b/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"] + diff --git a/packages/pynumaflow-lite/manifests/sink/README.md b/packages/pynumaflow-lite/manifests/sink/README.md new file mode 100644 index 00000000..ccfdd977 --- /dev/null +++ b/packages/pynumaflow-lite/manifests/sink/README.md @@ -0,0 +1,46 @@ +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 +``` + +Load it now to `k3d` + +```bash +k3d image import quay.io/numaio/numaflow/pynumaflow-lite-sink-log:v1 +``` + +## Run the pipeline + +```bash +kubectl apply -f pipeline.yaml +``` + +## Example: Using different response types + +The sink implementation supports all 5 response types: + +```python +# Success - write was successful +responses.append(sinker.Response.as_success(msg.id)) + +# Failure - write failed with error message +responses.append(sinker.Response.as_failure(msg.id, "Database connection failed")) + +# Fallback - forward to fallback sink +responses.append(sinker.Response.as_fallback(msg.id)) + +# Serve - write to serving store with payload +responses.append(sinker.Response.as_serve(msg.id, b"serving data")) + +# OnSuccess - forward to onSuccess store with optional message +message = sinker.Message( + value=b"processed data", + keys=["key1", "key2"], + user_metadata={"meta": sinker.KeyValueGroup.from_dict({"k": b"v"})} +) +responses.append(sinker.Response.as_on_success(msg.id, message)) +``` + diff --git a/packages/pynumaflow-lite/manifests/sink/pipeline.yaml b/packages/pynumaflow-lite/manifests/sink/pipeline.yaml new file mode 100644 index 00000000..b6ec043f --- /dev/null +++ b/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 + diff --git a/packages/pynumaflow-lite/manifests/sink/pyproject.toml b/packages/pynumaflow-lite/manifests/sink/pyproject.toml new file mode 100644 index 00000000..8e475beb --- /dev/null +++ b/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 = "vigith@gmail.com" } +] +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" + diff --git a/packages/pynumaflow-lite/manifests/sink/sink_log.py b/packages/pynumaflow-lite/manifests/sink/sink_log.py new file mode 100644 index 00000000..ad41b36a --- /dev/null +++ b/packages/pynumaflow-lite/manifests/sink/sink_log.py @@ -0,0 +1,62 @@ +import asyncio +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: 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)) + From 7c73887865c16fb2d931bf97be4e58001a21df5f Mon Sep 17 00:00:00 2001 From: Vigith Maurice Date: Sat, 8 Nov 2025 11:33:45 -0800 Subject: [PATCH 6/7] chore: clean up Signed-off-by: Vigith Maurice --- .../pynumaflow-lite/manifests/sink/README.md | 32 ++++--------------- 1 file changed, 7 insertions(+), 25 deletions(-) diff --git a/packages/pynumaflow-lite/manifests/sink/README.md b/packages/pynumaflow-lite/manifests/sink/README.md index ccfdd977..ed5c96d3 100644 --- a/packages/pynumaflow-lite/manifests/sink/README.md +++ b/packages/pynumaflow-lite/manifests/sink/README.md @@ -6,41 +6,23 @@ To create the `wheel` file, refer [root](../../README.md) 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 ``` -## Run the pipeline +### Minikube ```bash -kubectl apply -f pipeline.yaml +minikube image load quay.io/numaio/numaflow/pynumaflow-lite-sink-log:v1 ``` -## Example: Using different response types - -The sink implementation supports all 5 response types: - -```python -# Success - write was successful -responses.append(sinker.Response.as_success(msg.id)) - -# Failure - write failed with error message -responses.append(sinker.Response.as_failure(msg.id, "Database connection failed")) - -# Fallback - forward to fallback sink -responses.append(sinker.Response.as_fallback(msg.id)) - -# Serve - write to serving store with payload -responses.append(sinker.Response.as_serve(msg.id, b"serving data")) +## Run the pipeline -# OnSuccess - forward to onSuccess store with optional message -message = sinker.Message( - value=b"processed data", - keys=["key1", "key2"], - user_metadata={"meta": sinker.KeyValueGroup.from_dict({"k": b"v"})} -) -responses.append(sinker.Response.as_on_success(msg.id, message)) +```bash +kubectl apply -f pipeline.yaml ``` From d09310ead7f70b6778a367270a64b5826705eba6 Mon Sep 17 00:00:00 2001 From: Vigith Maurice Date: Mon, 10 Nov 2025 09:50:44 -0800 Subject: [PATCH 7/7] chore: code review Signed-off-by: Vigith Maurice --- packages/pynumaflow-lite/manifests/sink/sink_log.py | 3 ++- packages/pynumaflow-lite/src/sink/server.rs | 5 ++--- 2 files changed, 4 insertions(+), 4 deletions(-) diff --git a/packages/pynumaflow-lite/manifests/sink/sink_log.py b/packages/pynumaflow-lite/manifests/sink/sink_log.py index ad41b36a..986b270e 100644 --- a/packages/pynumaflow-lite/manifests/sink/sink_log.py +++ b/packages/pynumaflow-lite/manifests/sink/sink_log.py @@ -1,4 +1,5 @@ import asyncio +import collections import logging import signal from collections.abc import AsyncIterable @@ -34,7 +35,7 @@ async def handler(self, datums: AsyncIterable[sinker.Datum]) -> sinker.Responses pass -async def start(f: callable): +async def start(f: collections.abc.Callable): server = sinker.SinkAsyncServer() # Register loop-level signal handlers so we control shutdown and avoid asyncio.run diff --git a/packages/pynumaflow-lite/src/sink/server.rs b/packages/pynumaflow-lite/src/sink/server.rs index 937eb9f8..d971d514 100644 --- a/packages/pynumaflow-lite/src/sink/server.rs +++ b/packages/pynumaflow-lite/src/sink/server.rs @@ -1,5 +1,5 @@ -use numaflow::sink; use numaflow::shared::ServerExtras; +use numaflow::sink; use pyo3::prelude::*; use std::sync::Arc; @@ -98,9 +98,8 @@ pub(super) async fn start( // if not finished, abort it if !sig_handle.is_finished() { println!("Aborting signal handler"); - let _ = sig_handle.abort(); + sig_handle.abort(); } result } -