From 3535373b7e5e04b0bd65df1864f522774c19585e Mon Sep 17 00:00:00 2001 From: Vigith Maurice Date: Tue, 11 Nov 2025 10:48:40 -0800 Subject: [PATCH 1/2] feat: source using rust sdk Signed-off-by: Vigith Maurice --- packages/pynumaflow-lite/Cargo.toml | 4 + .../pynumaflow_lite/__init__.py | 28 +- .../pynumaflow_lite/__init__.pyi | 3 +- .../pynumaflow_lite/_source_dtypes.py | 125 +++++++ .../pynumaflow_lite/sourcer.pyi | 135 +++++++ packages/pynumaflow-lite/src/lib.rs | 18 + packages/pynumaflow-lite/src/source/mod.rs | 344 ++++++++++++++++++ packages/pynumaflow-lite/src/source/server.rs | 228 ++++++++++++ packages/pynumaflow-lite/tests/bin/source.rs | 214 +++++++++++ .../tests/examples/source_simple.py | 117 ++++++ packages/pynumaflow-lite/tests/test_source.py | 23 ++ 11 files changed, 1236 insertions(+), 3 deletions(-) create mode 100644 packages/pynumaflow-lite/pynumaflow_lite/_source_dtypes.py create mode 100644 packages/pynumaflow-lite/pynumaflow_lite/sourcer.pyi create mode 100644 packages/pynumaflow-lite/src/source/mod.rs create mode 100644 packages/pynumaflow-lite/src/source/server.rs create mode 100644 packages/pynumaflow-lite/tests/bin/source.rs create mode 100644 packages/pynumaflow-lite/tests/examples/source_simple.py create mode 100644 packages/pynumaflow-lite/tests/test_source.py diff --git a/packages/pynumaflow-lite/Cargo.toml b/packages/pynumaflow-lite/Cargo.toml index 5182f0e6..50824570 100644 --- a/packages/pynumaflow-lite/Cargo.toml +++ b/packages/pynumaflow-lite/Cargo.toml @@ -51,3 +51,7 @@ path = "tests/bin/accumulator.rs" [[bin]] name = "test_sink" path = "tests/bin/sink.rs" + +[[bin]] +name = "test_source" +path = "tests/bin/source.rs" diff --git a/packages/pynumaflow-lite/pynumaflow_lite/__init__.py b/packages/pynumaflow-lite/pynumaflow_lite/__init__.py index 680d066b..da65d338 100644 --- a/packages/pynumaflow-lite/pynumaflow_lite/__init__.py +++ b/packages/pynumaflow-lite/pynumaflow_lite/__init__.py @@ -33,13 +33,25 @@ except Exception: # pragma: no cover accumulator = None -# Surface the Python Mapper, BatchMapper, MapStreamer, Reducer, SessionReducer, and Accumulator classes under the extension submodules for convenient access +try: + sinker = _import_module(__name__ + ".sinker") +except Exception: # pragma: no cover + sinker = None + +try: + sourcer = _import_module(__name__ + ".sourcer") +except Exception: # pragma: no cover + sourcer = None + +# Surface the Python Mapper, BatchMapper, MapStreamer, Reducer, SessionReducer, Accumulator, Sinker, and Sourcer classes under the extension submodules for convenient access from ._map_dtypes import Mapper from ._batchmapper_dtypes import BatchMapper from ._mapstream_dtypes import MapStreamer from ._reduce_dtypes import Reducer from ._session_reduce_dtypes import SessionReducer from ._accumulator_dtypes import Accumulator +from ._sink_dtypes import Sinker +from ._source_dtypes import Sourcer if mapper is not None: try: @@ -77,8 +89,20 @@ except Exception: pass +if sinker is not None: + try: + setattr(sinker, "Sinker", Sinker) + except Exception: + pass + +if sourcer is not None: + try: + setattr(sourcer, "Sourcer", Sourcer) + except Exception: + pass + # Public API -__all__ = ["mapper", "batchmapper", "mapstreamer", "reducer", "session_reducer", "accumulator"] +__all__ = ["mapper", "batchmapper", "mapstreamer", "reducer", "session_reducer", "accumulator", "sinker", "sourcer"] __doc__ = pynumaflow_lite.__doc__ if hasattr(pynumaflow_lite, "__all__"): diff --git a/packages/pynumaflow-lite/pynumaflow_lite/__init__.pyi b/packages/pynumaflow-lite/pynumaflow_lite/__init__.pyi index d7e03a7b..6d6cbbb7 100644 --- a/packages/pynumaflow-lite/pynumaflow_lite/__init__.pyi +++ b/packages/pynumaflow-lite/pynumaflow_lite/__init__.pyi @@ -9,5 +9,6 @@ from . import reducer as reducer from . import session_reducer as session_reducer from . import accumulator as accumulator from . import sinker as sinker +from . import sourcer as sourcer -__all__ = ['mapper', 'batchmapper', 'mapstreamer', 'reducer', 'session_reducer', 'accumulator', 'sinker'] +__all__ = ['mapper', 'batchmapper', 'mapstreamer', 'reducer', 'session_reducer', 'accumulator', 'sinker', 'sourcer'] diff --git a/packages/pynumaflow-lite/pynumaflow_lite/_source_dtypes.py b/packages/pynumaflow-lite/pynumaflow_lite/_source_dtypes.py new file mode 100644 index 00000000..1ab72137 --- /dev/null +++ b/packages/pynumaflow-lite/pynumaflow_lite/_source_dtypes.py @@ -0,0 +1,125 @@ +from abc import ABCMeta, abstractmethod +from collections.abc import AsyncIterator +from pynumaflow_lite.sourcer import ( + Message, + Offset, + ReadRequest, + AckRequest, + NackRequest, + PendingResponse, + PartitionsResponse, +) + + +class Sourcer(metaclass=ABCMeta): + """ + Provides an interface to write a User Defined Source. + + A Sourcer must implement the following handlers: + - read_handler: Read messages from the source + - ack_handler: Acknowledge processed messages + - pending_handler: Return the number of pending messages + - partitions_handler: Return the partitions this source handles + + Optionally, you can implement: + - nack_handler: Negatively acknowledge messages (default: no-op) + """ + + def __call__(self, *args, **kwargs): + """ + This allows to execute the handler function directly if + class instance is sent as a callable. + """ + return self.read_handler(*args, **kwargs) + + @abstractmethod + async def read_handler(self, request: ReadRequest) -> AsyncIterator[Message]: + """ + Read messages from the source. + + Args: + request: ReadRequest containing num_records and timeout + + Yields: + Message: Messages to be sent to the next vertex + + Example: + async def read_handler(self, request: ReadRequest) -> AsyncIterator[Message]: + for i in range(request.num_records): + yield Message( + payload=f"message-{i}".encode(), + offset=Offset(str(i).encode(), partition_id=0), + event_time=datetime.now(), + keys=["key1"], + headers={"x-txn-id": str(uuid.uuid4())} + ) + """ + pass + + @abstractmethod + async def ack_handler(self, request: AckRequest) -> None: + """ + Acknowledge that messages have been processed. + + Args: + request: AckRequest containing the list of offsets to acknowledge + + Example: + async def ack_handler(self, request: AckRequest) -> None: + for offset in request.offsets: + # Remove from pending set, mark as processed, etc. + self.pending_offsets.remove(offset.offset) + """ + pass + + @abstractmethod + async def pending_handler(self) -> PendingResponse: + """ + Return the number of pending messages yet to be processed. + + Returns: + PendingResponse: Response containing the count of pending messages. + Return count=-1 if the source doesn't support detecting backlog. + + Example: + async def pending_handler(self) -> PendingResponse: + return PendingResponse(count=len(self.pending_offsets)) + """ + pass + + @abstractmethod + async def partitions_handler(self) -> PartitionsResponse: + """ + Return the partitions associated with this source. + + This is used by the platform to determine the partitions to which + the watermark should be published. If your source doesn't have the + concept of partitions, return the replica ID. + + Returns: + PartitionsResponse: Response containing the list of partition IDs + + Example: + async def partitions_handler(self) -> PartitionsResponse: + return PartitionsResponse(partitions=[self.partition_id]) + """ + pass + + async def nack_handler(self, request: NackRequest) -> None: + """ + Negatively acknowledge messages (optional). + + This is called when messages could not be processed and should be + retried or handled differently. Default implementation is a no-op. + + Args: + request: NackRequest containing the list of offsets to nack + + Example: + async def nack_handler(self, request: NackRequest) -> None: + for offset in request.offsets: + # Add back to pending, mark for retry, etc. + self.nacked_offsets.add(offset.offset) + """ + pass + diff --git a/packages/pynumaflow-lite/pynumaflow_lite/sourcer.pyi b/packages/pynumaflow-lite/pynumaflow_lite/sourcer.pyi new file mode 100644 index 00000000..1b0b1e84 --- /dev/null +++ b/packages/pynumaflow-lite/pynumaflow_lite/sourcer.pyi @@ -0,0 +1,135 @@ +from __future__ import annotations + +from typing import Optional, List, Dict, Callable, Awaitable, Any +import datetime as _dt + +# Re-export the Python ABC for user convenience and typing +from ._source_dtypes import Sourcer as Sourcer + + +class Message: + """A message to be sent from the source.""" + payload: bytes + offset: Offset + event_time: _dt.datetime + keys: List[str] + headers: Dict[str, str] + + def __init__( + self, + payload: bytes, + offset: Offset, + event_time: _dt.datetime, + keys: Optional[List[str]] = ..., + headers: Optional[Dict[str, str]] = ..., + ) -> None: ... + + def __repr__(self) -> str: ... + + def __str__(self) -> str: ... + + +class Offset: + """The offset of a message.""" + offset: bytes + partition_id: int + + def __init__( + self, + offset: bytes, + partition_id: int = ..., + ) -> None: ... + + def __repr__(self) -> str: ... + + def __str__(self) -> str: ... + + +class ReadRequest: + """A request to read messages from the source.""" + num_records: int + timeout_ms: int + + def __init__( + self, + num_records: int, + timeout_ms: int = ..., + ) -> None: ... + + def __repr__(self) -> str: ... + + +class AckRequest: + """A request to acknowledge messages.""" + offsets: List[Offset] + + def __init__( + self, + offsets: List[Offset], + ) -> None: ... + + def __repr__(self) -> str: ... + + +class NackRequest: + """A request to negatively acknowledge messages.""" + offsets: List[Offset] + + def __init__( + self, + offsets: List[Offset], + ) -> None: ... + + def __repr__(self) -> str: ... + + +class PendingResponse: + """Response for pending messages count.""" + count: int + + def __init__( + self, + count: int = ..., + ) -> None: ... + + def __repr__(self) -> str: ... + + +class PartitionsResponse: + """Response for partitions.""" + partitions: List[int] + + def __init__( + self, + partitions: List[int], + ) -> None: ... + + def __repr__(self) -> str: ... + + +class SourceAsyncServer: + """Async Source Server that can be started from Python code.""" + + 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: ... + + +__all__ = [ + "Message", + "Offset", + "ReadRequest", + "AckRequest", + "NackRequest", + "PendingResponse", + "PartitionsResponse", + "SourceAsyncServer", + "Sourcer", +] + diff --git a/packages/pynumaflow-lite/src/lib.rs b/packages/pynumaflow-lite/src/lib.rs index c5f7af21..200b3813 100644 --- a/packages/pynumaflow-lite/src/lib.rs +++ b/packages/pynumaflow-lite/src/lib.rs @@ -7,6 +7,7 @@ pub mod pyrs; pub mod reduce; pub mod session_reduce; pub mod sink; +pub mod source; use pyo3::prelude::*; @@ -59,6 +60,13 @@ fn sinker(_py: Python, m: &Bound) -> PyResult<()> { Ok(()) } +/// Submodule: pynumaflow_lite.sourcer +#[pymodule] +fn sourcer(_py: Python, m: &Bound) -> PyResult<()> { + crate::source::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<()> { @@ -70,6 +78,7 @@ fn pynumaflow_lite(py: Python, m: &Bound) -> PyResult<()> { m.add_wrapped(pyo3::wrap_pymodule!(session_reducer))?; m.add_wrapped(pyo3::wrap_pymodule!(accumulator))?; m.add_wrapped(pyo3::wrap_pymodule!(sinker))?; + m.add_wrapped(pyo3::wrap_pymodule!(sourcer))?; // Ensure it's importable as `pynumaflow_lite.mapper` as well as attribute access let binding = m.getattr("mapper")?; @@ -134,5 +143,14 @@ fn pynumaflow_lite(py: Python, m: &Bound) -> PyResult<()> { .getattr("modules")? .set_item(fullname, &sub)?; + // Ensure it's importable as `pynumaflow_lite.sourcer` as well + let binding = m.getattr("sourcer")?; + let sub = binding.downcast::()?; + let fullname = "pynumaflow_lite.sourcer"; + sub.setattr("__name__", fullname)?; + py.import("sys")? + .getattr("modules")? + .set_item(fullname, &sub)?; + Ok(()) } diff --git a/packages/pynumaflow-lite/src/source/mod.rs b/packages/pynumaflow-lite/src/source/mod.rs new file mode 100644 index 00000000..ed6b4b8d --- /dev/null +++ b/packages/pynumaflow-lite/src/source/mod.rs @@ -0,0 +1,344 @@ +use std::collections::HashMap; + +use chrono::{DateTime, Utc}; + +/// Source interface managed by Python. It means Python code will start the server +/// and can pass in the Python function. +pub mod server; + +use pyo3::prelude::*; +use std::sync::Mutex; + +/// A message to be sent from the source. +#[pyclass(module = "pynumaflow_lite.sourcer")] +#[derive(Clone, Debug)] +pub struct Message { + /// The payload of the message. + #[pyo3(get)] + pub payload: Vec, + /// The offset of the message (not directly exposed, use offset property). + pub(crate) offset: PyOffset, + /// The event time of the message. + #[pyo3(get)] + pub event_time: DateTime, + /// Keys of the message. + #[pyo3(get)] + pub keys: Vec, + /// Headers of the message. + #[pyo3(get)] + pub headers: HashMap, +} + +#[pymethods] +impl Message { + /// Create a new [Message] with the given payload, offset, event_time, keys, and headers. + #[new] + #[pyo3(signature = (payload: "bytes", offset: "Offset", event_time: "datetime", keys: "list[str] | None"=None, headers: "dict[str, str] | None"=None) -> "Message" + )] + fn new( + payload: Vec, + offset: PyOffset, + event_time: DateTime, + keys: Option>, + headers: Option>, + ) -> Self { + Self { + payload, + offset, + event_time, + keys: keys.unwrap_or_default(), + headers: headers.unwrap_or_default(), + } + } + + /// Get the offset of the message. + #[getter] + fn offset(&self) -> PyOffset { + self.offset.clone() + } + + fn __repr__(&self) -> String { + format!( + "Message(payload={:?}, offset={:?}, event_time={}, keys={:?}, headers={:?})", + self.payload, self.offset, self.event_time, self.keys, self.headers + ) + } + + fn __str__(&self) -> String { + format!( + "Message(payload={:?}, offset={:?}, event_time={}, keys={:?}, headers={:?})", + String::from_utf8_lossy(&self.payload), + self.offset, + self.event_time, + self.keys, + self.headers + ) + } +} + +impl From for numaflow::source::Message { + fn from(value: Message) -> Self { + Self { + value: value.payload, + offset: value.offset.into(), + event_time: value.event_time, + keys: value.keys, + headers: value.headers, + } + } +} + +/// The offset of a message. +#[pyclass(module = "pynumaflow_lite.sourcer", name = "Offset")] +#[derive(Clone, Debug)] +pub struct PyOffset { + /// Offset value in bytes. + #[pyo3(get)] + pub offset: Vec, + /// Partition ID of the message. + #[pyo3(get)] + pub partition_id: i32, +} + +#[pymethods] +impl PyOffset { + /// Create a new [Offset] with the given offset bytes and partition_id. + #[new] + #[pyo3(signature = (offset: "bytes", partition_id: "int"=0) -> "Offset")] + fn new(offset: Vec, partition_id: i32) -> Self { + Self { + offset, + partition_id, + } + } + + fn __repr__(&self) -> String { + format!( + "Offset(offset={:?}, partition_id={})", + self.offset, self.partition_id + ) + } + + fn __str__(&self) -> String { + format!( + "Offset(offset={:?}, partition_id={})", + String::from_utf8_lossy(&self.offset), + self.partition_id + ) + } +} + +impl From for numaflow::source::Offset { + fn from(value: PyOffset) -> Self { + Self { + offset: value.offset, + partition_id: value.partition_id, + } + } +} + +impl From for PyOffset { + fn from(value: numaflow::source::Offset) -> Self { + Self { + offset: value.offset, + partition_id: value.partition_id, + } + } +} + +/// A request to read messages from the source. +#[pyclass(module = "pynumaflow_lite.sourcer")] +#[derive(Clone, Debug)] +pub struct ReadRequest { + /// The number of messages to read. + #[pyo3(get)] + pub num_records: u64, + /// Request timeout in milliseconds. + #[pyo3(get)] + pub timeout_ms: u64, +} + +#[pymethods] +impl ReadRequest { + /// Create a new [ReadRequest] with the given num_records and timeout_ms. + #[new] + #[pyo3(signature = (num_records: "int", timeout_ms: "int"=1000) -> "ReadRequest")] + fn new(num_records: u64, timeout_ms: u64) -> Self { + Self { + num_records, + timeout_ms, + } + } + + fn __repr__(&self) -> String { + format!( + "ReadRequest(num_records={}, timeout_ms={})", + self.num_records, self.timeout_ms + ) + } +} + +impl From for ReadRequest { + fn from(value: numaflow::source::SourceReadRequest) -> Self { + Self { + num_records: value.count as u64, + timeout_ms: value.timeout.as_millis() as u64, + } + } +} + +/// A request to acknowledge messages. +#[pyclass(module = "pynumaflow_lite.sourcer")] +#[derive(Clone, Debug)] +pub struct AckRequest { + /// The offsets to acknowledge. + #[pyo3(get)] + pub offsets: Vec, +} + +#[pymethods] +impl AckRequest { + /// Create a new [AckRequest] with the given offsets. + #[new] + #[pyo3(signature = (offsets: "list[Offset]") -> "AckRequest")] + fn new(offsets: Vec) -> Self { + Self { offsets } + } + + fn __repr__(&self) -> String { + format!("AckRequest(offsets={:?})", self.offsets) + } +} + +/// A request to negatively acknowledge messages. +#[pyclass(module = "pynumaflow_lite.sourcer")] +#[derive(Clone, Debug)] +pub struct NackRequest { + /// The offsets to negatively acknowledge. + #[pyo3(get)] + pub offsets: Vec, +} + +#[pymethods] +impl NackRequest { + /// Create a new [NackRequest] with the given offsets. + #[new] + #[pyo3(signature = (offsets: "list[Offset]") -> "NackRequest")] + fn new(offsets: Vec) -> Self { + Self { offsets } + } + + fn __repr__(&self) -> String { + format!("NackRequest(offsets={:?})", self.offsets) + } +} + +/// Response for pending messages count. +#[pyclass(module = "pynumaflow_lite.sourcer")] +#[derive(Clone, Debug)] +pub struct PendingResponse { + /// The number of pending messages. -1 if the source doesn't support detecting backlog. + #[pyo3(get)] + pub count: i64, +} + +#[pymethods] +impl PendingResponse { + /// Create a new [PendingResponse] with the given count. + #[new] + #[pyo3(signature = (count: "int"=0) -> "PendingResponse")] + fn new(count: i64) -> Self { + Self { count } + } + + fn __repr__(&self) -> String { + format!("PendingResponse(count={})", self.count) + } +} + +/// Response for partitions. +#[pyclass(module = "pynumaflow_lite.sourcer")] +#[derive(Clone, Debug)] +pub struct PartitionsResponse { + /// The list of partition IDs. + #[pyo3(get)] + pub partitions: Vec, +} + +#[pymethods] +impl PartitionsResponse { + /// Create a new [PartitionsResponse] with the given partitions. + #[new] + #[pyo3(signature = (partitions: "list[int]") -> "PartitionsResponse")] + fn new(partitions: Vec) -> Self { + Self { partitions } + } + + fn __repr__(&self) -> String { + format!("PartitionsResponse(partitions={:?})", self.partitions) + } +} + +/// Async Source Server that can be started from Python code which will run the Python UDF function. +#[pyclass(module = "pynumaflow_lite.sourcer")] +pub struct SourceAsyncServer { + sock_file: String, + info_file: String, + shutdown_tx: Mutex>>, +} + +#[pymethods] +impl SourceAsyncServer { + #[new] + #[pyo3(signature = (sock_file: "str | None"=numaflow::source::SOCK_ADDR.to_string(), info_file: "str | None"=numaflow::source::SERVER_INFO_FILE.to_string()) -> "SourceAsyncServer" + )] + 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 source handler. + #[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::source::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 source 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::()?; + m.add_class::()?; + + Ok(()) +} diff --git a/packages/pynumaflow-lite/src/source/server.rs b/packages/pynumaflow-lite/src/source/server.rs new file mode 100644 index 00000000..e4c6a545 --- /dev/null +++ b/packages/pynumaflow-lite/src/source/server.rs @@ -0,0 +1,228 @@ +use numaflow::shared::ServerExtras; + +use pyo3::prelude::*; +use std::sync::Arc; +use tokio::sync::mpsc::Sender; + +pub(crate) struct PySourceRunner { + pub(crate) event_loop: Arc>, + pub(crate) py_handler: Arc>, +} + +#[tonic::async_trait] +impl numaflow::source::Sourcer for PySourceRunner { + /// Reads the messages from the source and sends them to the transmitter. + async fn read( + &self, + request: numaflow::source::SourceReadRequest, + transmitter: Sender, + ) { + // Convert the Rust request to Python ReadRequest + let read_request: crate::source::ReadRequest = request.into(); + + // Call the Python read_handler which returns an AsyncIterator[Message] + let py_async_iter = Python::attach(|py| { + let py_handler = self.py_handler.clone(); + + // Call read_handler(request) -> AsyncIterator[Message] + let result = py_handler + .call_method1(py, "read_handler", (read_request,)) + .expect("failed to call read_handler"); + + result + }); + + // Create a stream from the Python AsyncIterator + let mut stream = crate::pyiterables::PyAsyncIterStream::::new( + py_async_iter, + self.event_loop.clone(), + ) + .expect("failed to create stream from Python AsyncIterator"); + + // Stream messages from Python to the transmitter + use tokio_stream::StreamExt; + while let Some(result) = stream.next().await { + match result { + Ok(py_message) => { + let rust_message: numaflow::source::Message = py_message.into(); + if transmitter.send(rust_message).await.is_err() { + // Receiver dropped, stop reading + break; + } + } + Err(e) => { + eprintln!("Error reading from Python source: {:?}", e); + break; + } + } + } + } + + /// Acknowledges the message that has been processed by the user-defined source. + async fn ack(&self, offsets: Vec) { + // Convert Rust offsets to Python Offset objects + let py_offsets: Vec = + offsets.into_iter().map(|o| o.into()).collect(); + + // Create AckRequest + let ack_request = crate::source::AckRequest::new(py_offsets); + + // Call the Python ack_handler + let fut = Python::attach(|py| { + let py_handler = self.py_handler.clone(); + let locals = pyo3_async_runtimes::TaskLocals::new(self.event_loop.bind(py).clone()); + + let coro = py_handler + .call_method1(py, "ack_handler", (ack_request,)) + .expect("failed to call ack_handler") + .into_bound(py); + + pyo3_async_runtimes::into_future_with_locals(&locals, coro) + .expect("failed to convert ack_handler to future") + }); + + // Await the Python coroutine + let _ = fut.await; + } + + /// Negatively acknowledges the message that has been processed by the user-defined source. + async fn nack(&self, offsets: Vec) { + // Convert Rust offsets to Python Offset objects + let py_offsets: Vec = + offsets.into_iter().map(|o| o.into()).collect(); + + // Create NackRequest + let nack_request = crate::source::NackRequest::new(py_offsets); + + // Call the Python nack_handler + let fut = Python::attach(|py| { + let py_handler = self.py_handler.clone(); + let locals = pyo3_async_runtimes::TaskLocals::new(self.event_loop.bind(py).clone()); + + let coro = py_handler + .call_method1(py, "nack_handler", (nack_request,)) + .expect("failed to call nack_handler") + .into_bound(py); + + pyo3_async_runtimes::into_future_with_locals(&locals, coro) + .expect("failed to convert nack_handler to future") + }); + + // Await the Python coroutine + let _ = fut.await; + } + + /// Returns the number of messages that are yet to be processed by the user-defined source. + /// The None value can be returned if source doesn't support detecting the backlog. + async fn pending(&self) -> Option { + // Call the Python pending_handler + let fut = Python::attach(|py| { + let py_handler = self.py_handler.clone(); + let locals = pyo3_async_runtimes::TaskLocals::new(self.event_loop.bind(py).clone()); + + let coro = py_handler + .call_method0(py, "pending_handler") + .expect("failed to call pending_handler") + .into_bound(py); + + pyo3_async_runtimes::into_future_with_locals(&locals, coro) + .expect("failed to convert pending_handler to future") + }); + + // Await the Python coroutine and extract the result + let result = fut.await.expect("pending_handler failed"); + + let pending_response = Python::attach(|py| { + result + .extract::(py) + .expect("failed to extract PendingResponse") + }); + + // Convert count to Option + // -1 means the source doesn't support detecting backlog + if pending_response.count < 0 { + None + } else { + Some(pending_response.count as usize) + } + } + + /// Returns the partitions associated with the source. This will be used by the platform to determine + /// the partitions to which the watermark should be published. Some sources might not have the concept of partitions. + /// Kafka is an example of source where a reader can read from multiple partitions. + /// If None is returned, Numaflow replica-id will be returned as the partition. + async fn partitions(&self) -> Option> { + // Call the Python partitions_handler + let fut = Python::attach(|py| { + let py_handler = self.py_handler.clone(); + let locals = pyo3_async_runtimes::TaskLocals::new(self.event_loop.bind(py).clone()); + + let coro = py_handler + .call_method0(py, "partitions_handler") + .expect("failed to call partitions_handler") + .into_bound(py); + + pyo3_async_runtimes::into_future_with_locals(&locals, coro) + .expect("failed to convert partitions_handler to future") + }); + + // Await the Python coroutine and extract the result + let result = fut.await.expect("partitions_handler failed"); + + let partitions_response = Python::attach(|py| { + result + .extract::(py) + .expect("failed to extract PartitionsResponse") + }); + + Some(partitions_response.partitions) + } +} + +/// Start the source server by spinning up a dedicated Python asyncio loop and wiring shutdown. +pub(super) async fn start( + py_handler: 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_source_runner = PySourceRunner { + py_handler: Arc::new(py_handler), + event_loop: event_loop.clone(), + }; + + let server = numaflow::source::Server::new(py_source_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 Source 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 +} diff --git a/packages/pynumaflow-lite/tests/bin/source.rs b/packages/pynumaflow-lite/tests/bin/source.rs new file mode 100644 index 00000000..460fb4c1 --- /dev/null +++ b/packages/pynumaflow-lite/tests/bin/source.rs @@ -0,0 +1,214 @@ +use std::env; +use std::path::PathBuf; + +use numaflow::proto; +use numaflow::proto::source::source_client::SourceClient; +use tokio::net::UnixStream; +use tokio::sync::mpsc; +use tokio_stream::wrappers::ReceiverStream; +use tonic::Request; +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_SOURCE_SOCK").ok()) + .unwrap_or_else(|| "/tmp/var/run/numaflow/source.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 = SourceClient::new(channel); + + // Test 1: Read messages + println!("Testing read operation..."); + let messages = read_messages(&mut client, 5).await?; + println!("Read {} messages", messages.len()); + assert_eq!(messages.len(), 5, "Should read exactly 5 messages"); + + // Verify message content + for (i, message) in messages.iter().enumerate() { + println!( + "Message {}: payload={}, keys={:?}, partition_id={}", + i, + String::from_utf8_lossy(&message.payload), + message.keys, + message.offset.as_ref().unwrap().partition_id + ); + assert!(!message.payload.is_empty(), "Message should have payload"); + assert!(message.offset.is_some(), "Message should have offset"); + } + + // Test 2: Pending before ack + println!("\nTesting pending operation before ack..."); + let pending_response = client.pending_fn(Request::new(())).await?; + let pending_count = pending_response.into_inner().result.unwrap().count; + println!("Pending messages: {}", pending_count); + // The simple source returns 0 for pending + assert_eq!(pending_count, 0, "Simple source should return 0 pending"); + + // Test 3: Partitions + println!("\nTesting partitions operation..."); + let partitions_response = client.partitions_fn(Request::new(())).await?; + let partitions = partitions_response.into_inner().result.unwrap().partitions; + println!("Partitions: {:?}", partitions); + assert!(!partitions.is_empty(), "Should have at least one partition"); + + // Test 4: Ack messages + println!("\nTesting ack operation..."); + ack_messages(&mut client, &messages).await?; + println!("Successfully acknowledged {} messages", messages.len()); + + // Test 5: Read more messages and test nack + println!("\nTesting nack operation..."); + let nack_messages = read_messages(&mut client, 2).await?; + println!("Read {} messages for nack test", nack_messages.len()); + nack_messages_fn(&mut client, &nack_messages).await?; + println!("Successfully nacked {} messages", nack_messages.len()); + + println!("\nAll source tests passed!"); + + Ok(()) +} + +/// Read messages from the source with proper handshake +async fn read_messages( + client: &mut SourceClient, + num_records: u64, +) -> Result, Box> { + let (read_tx, read_rx) = mpsc::channel(4); + + // Send handshake + let handshake_request = proto::source::ReadRequest { + request: None, + handshake: Some(proto::source::Handshake { sot: true }), + }; + read_tx.send(handshake_request).await?; + + // Send read request + let read_request = proto::source::ReadRequest { + request: Some(proto::source::read_request::Request { + num_records, + timeout_in_ms: 1000, + }), + handshake: None, + }; + read_tx.send(read_request).await?; + drop(read_tx); + + let mut response_stream = client + .read_fn(Request::new(ReceiverStream::new(read_rx))) + .await? + .into_inner(); + + // Expect handshake response + let handshake_response = response_stream.message().await?.unwrap(); + assert!( + handshake_response.handshake.is_some(), + "Should receive handshake response" + ); + + let mut messages = Vec::new(); + while let Some(response) = response_stream.message().await? { + if let Some(status) = response.status { + if status.eot { + break; + } + } + if let Some(result) = response.result { + messages.push(result); + } + } + + Ok(messages) +} + +/// Acknowledge messages with proper handshake +async fn ack_messages( + client: &mut SourceClient, + messages: &[proto::source::read_response::Result], +) -> Result<(), Box> { + let (ack_tx, ack_rx) = mpsc::channel(10); + + // Send handshake + let ack_handshake_request = proto::source::AckRequest { + request: None, + handshake: Some(proto::source::Handshake { sot: true }), + }; + ack_tx.send(ack_handshake_request).await?; + + // Send ack requests + for message in messages { + let ack_request = proto::source::AckRequest { + request: Some(proto::source::ack_request::Request { + offsets: vec![proto::source::Offset { + offset: message.offset.as_ref().unwrap().offset.clone(), + partition_id: message.offset.as_ref().unwrap().partition_id, + }], + }), + handshake: None, + }; + ack_tx.send(ack_request).await?; + } + drop(ack_tx); + + let mut ack_response = client + .ack_fn(Request::new(ReceiverStream::new(ack_rx))) + .await? + .into_inner(); + + // Consume handshake response + let handshake_response = ack_response.message().await?.unwrap(); + assert!( + handshake_response.handshake.unwrap().sot, + "Should receive ack handshake" + ); + + // Consume ack responses + for _ in 0..messages.len() { + let ack_result = ack_response.message().await?.unwrap(); + assert!( + ack_result.result.is_some(), + "Should receive ack result for each message" + ); + } + + Ok(()) +} + +/// Negatively acknowledge messages +async fn nack_messages_fn( + client: &mut SourceClient, + messages: &[proto::source::read_response::Result], +) -> Result<(), Box> { + for message in messages { + let nack_request = proto::source::NackRequest { + request: Some(proto::source::nack_request::Request { + offsets: vec![proto::source::Offset { + offset: message.offset.as_ref().unwrap().offset.clone(), + partition_id: message.offset.as_ref().unwrap().partition_id, + }], + }), + }; + + let nack_response = client.nack_fn(Request::new(nack_request)).await?; + assert!( + nack_response.into_inner().result.is_some(), + "Should receive nack result" + ); + } + + Ok(()) +} diff --git a/packages/pynumaflow-lite/tests/examples/source_simple.py b/packages/pynumaflow-lite/tests/examples/source_simple.py new file mode 100644 index 00000000..2517e6a0 --- /dev/null +++ b/packages/pynumaflow-lite/tests/examples/source_simple.py @@ -0,0 +1,117 @@ +import asyncio +import logging +import signal +from datetime import datetime, timezone +from collections.abc import AsyncIterator + +from pynumaflow_lite import sourcer +from pynumaflow_lite._source_dtypes import Sourcer + +# Configure logging +logging.basicConfig(level=logging.INFO) +_LOGGER = logging.getLogger(__name__) + + +class SimpleSource(Sourcer): + """ + Simple source that generates messages with incrementing numbers. + This is the class-based approach matching the user's example. + """ + + def __init__(self): + self.counter = 0 + self.partition_idx = 0 + + async def read_handler(self, datum: sourcer.ReadRequest) -> AsyncIterator[sourcer.Message]: + """ + The simple source generates messages with incrementing numbers. + """ + _LOGGER.info(f"Read request: num_records={datum.num_records}, timeout_ms={datum.timeout_ms}") + + # Generate the requested number of messages + for i in range(datum.num_records): + # Create message payload + payload = f"message-{self.counter}".encode("utf-8") + + # Create offset + offset = sourcer.Offset( + offset=str(self.counter).encode("utf-8"), + partition_id=self.partition_idx + ) + + # Create message + message = sourcer.Message( + payload=payload, + offset=offset, + event_time=datetime.now(timezone.utc), + keys=["key1"], + headers={"source": "simple"} + ) + + _LOGGER.info(f"Generated message: {self.counter}") + self.counter += 1 + + yield message + + # Small delay to simulate real source + await asyncio.sleep(0.1) + + async def ack_handler(self, request: sourcer.AckRequest) -> None: + """ + The simple source acknowledges the offsets. + """ + _LOGGER.info(f"Acknowledging {len(request.offsets)} offsets") + for offset in request.offsets: + _LOGGER.debug(f"Acked offset: {offset.offset.decode('utf-8')}, partition: {offset.partition_id}") + + async def nack_handler(self, request: sourcer.NackRequest) -> None: + """ + The simple source negatively acknowledges the offsets. + """ + _LOGGER.info(f"Negatively acknowledging {len(request.offsets)} offsets") + for offset in request.offsets: + _LOGGER.warning(f"Nacked offset: {offset.offset.decode('utf-8')}, partition: {offset.partition_id}") + + async def pending_handler(self) -> sourcer.PendingResponse: + """ + The simple source always returns zero to indicate there is no pending record. + """ + return sourcer.PendingResponse(count=0) + + async def partitions_handler(self) -> sourcer.PartitionsResponse: + """ + The simple source always returns default partitions. + """ + return sourcer.PartitionsResponse(partitions=[self.partition_idx]) + + +async def start(): + sock_file = "/tmp/var/run/numaflow/source.sock" + server_info_file = "/tmp/var/run/numaflow/sourcer-server-info" + server = sourcer.SourceAsyncServer(sock_file, server_info_file) + + # Create an instance of the source handler + handler = SimpleSource() + + # 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_source.py b/packages/pynumaflow-lite/tests/test_source.py new file mode 100644 index 00000000..beb40d12 --- /dev/null +++ b/packages/pynumaflow-lite/tests/test_source.py @@ -0,0 +1,23 @@ +from pathlib import Path + +import pytest + +from _test_utils import run_python_server_with_rust_client + +SOCK_PATH = Path("/tmp/var/run/numaflow/source.sock") +SERVER_INFO = Path("/tmp/var/run/numaflow/sourcer-server-info") + +SCRIPTS = [ + "source_simple.py", +] + + +@pytest.mark.parametrize("script", SCRIPTS) +def test_python_source_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_source", + ) + From cd0a4df6186795632e66f6214010880af6d62c42 Mon Sep 17 00:00:00 2001 From: Vigith Maurice Date: Tue, 11 Nov 2025 11:05:49 -0800 Subject: [PATCH 2/2] chore: tidy Signed-off-by: Vigith Maurice --- packages/pynumaflow-lite/src/sink/mod.rs | 1 - packages/pynumaflow-lite/src/source/mod.rs | 2 +- packages/pynumaflow-lite/src/source/server.rs | 2 +- 3 files changed, 2 insertions(+), 3 deletions(-) diff --git a/packages/pynumaflow-lite/src/sink/mod.rs b/packages/pynumaflow-lite/src/sink/mod.rs index 6ccac8e7..15087090 100644 --- a/packages/pynumaflow-lite/src/sink/mod.rs +++ b/packages/pynumaflow-lite/src/sink/mod.rs @@ -398,4 +398,3 @@ pub(crate) fn populate_py_module(m: &Bound) -> PyResult<()> { Ok(()) } - diff --git a/packages/pynumaflow-lite/src/source/mod.rs b/packages/pynumaflow-lite/src/source/mod.rs index ed6b4b8d..3fcd7a7a 100644 --- a/packages/pynumaflow-lite/src/source/mod.rs +++ b/packages/pynumaflow-lite/src/source/mod.rs @@ -89,7 +89,7 @@ impl From for numaflow::source::Message { } /// The offset of a message. -#[pyclass(module = "pynumaflow_lite.sourcer", name = "Offset")] +#[pyclass(module = "pynumaflow_lite.sourcer", name = "Offset")] // this to avoid conflict with the Offset in the source module #[derive(Clone, Debug)] pub struct PyOffset { /// Offset value in bytes. diff --git a/packages/pynumaflow-lite/src/source/server.rs b/packages/pynumaflow-lite/src/source/server.rs index e4c6a545..f47540e9 100644 --- a/packages/pynumaflow-lite/src/source/server.rs +++ b/packages/pynumaflow-lite/src/source/server.rs @@ -221,7 +221,7 @@ 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