diff --git a/packages/pynumaflow-lite/Cargo.toml b/packages/pynumaflow-lite/Cargo.toml index 3c0939ac..c5025cf0 100644 --- a/packages/pynumaflow-lite/Cargo.toml +++ b/packages/pynumaflow-lite/Cargo.toml @@ -59,3 +59,7 @@ path = "tests/bin/sink.rs" [[bin]] name = "test_source" path = "tests/bin/source.rs" + +[[bin]] +name = "test_sourcetransform" +path = "tests/bin/sourcetransform.rs" diff --git a/packages/pynumaflow-lite/pynumaflow_lite/__init__.py b/packages/pynumaflow-lite/pynumaflow_lite/__init__.py index f29648be..07b309fe 100644 --- a/packages/pynumaflow-lite/pynumaflow_lite/__init__.py +++ b/packages/pynumaflow-lite/pynumaflow_lite/__init__.py @@ -48,7 +48,13 @@ except Exception: # pragma: no cover sourcer = None -# Surface the Python Mapper, BatchMapper, MapStreamer, Reducer, SessionReducer, ReduceStreamer, Accumulator, Sinker, and Sourcer classes under the extension submodules for convenient access +try: + sourcetransformer = _import_module(__name__ + ".sourcetransformer") +except Exception: # pragma: no cover + sourcetransformer = None + +# Surface the Python Mapper, BatchMapper, MapStreamer, Reducer, SessionReducer, ReduceStreamer, Accumulator, Sinker, +# Sourcer, and SourceTransformer classes under the extension submodules for convenient access from ._map_dtypes import Mapper from ._batchmapper_dtypes import BatchMapper from ._mapstream_dtypes import MapStreamer @@ -58,6 +64,7 @@ from ._accumulator_dtypes import Accumulator from ._sink_dtypes import Sinker from ._source_dtypes import Sourcer +from ._sourcetransformer_dtypes import SourceTransformer if mapper is not None: try: @@ -113,8 +120,15 @@ except Exception: pass +if sourcetransformer is not None: + try: + setattr(sourcetransformer, "SourceTransformer", SourceTransformer) + except Exception: + pass + # Public API -__all__ = ["mapper", "batchmapper", "mapstreamer", "reducer", "session_reducer", "reducestreamer", "accumulator", "sinker", "sourcer"] +__all__ = ["mapper", "batchmapper", "mapstreamer", "reducer", "session_reducer", "reducestreamer", "accumulator", + "sinker", "sourcer", "sourcetransformer"] __doc__ = pynumaflow_lite.__doc__ if hasattr(pynumaflow_lite, "__all__"): diff --git a/packages/pynumaflow-lite/pynumaflow_lite/_sourcetransformer_dtypes.py b/packages/pynumaflow-lite/pynumaflow_lite/_sourcetransformer_dtypes.py new file mode 100644 index 00000000..0727b3b3 --- /dev/null +++ b/packages/pynumaflow-lite/pynumaflow_lite/_sourcetransformer_dtypes.py @@ -0,0 +1,35 @@ +from abc import ABCMeta, abstractmethod +from pynumaflow_lite.sourcetransformer import Datum, Messages + + +class SourceTransformer(metaclass=ABCMeta): + """ + Provides an interface to write a SourceTransformer + which will be exposed over a gRPC server. + + A SourceTransformer is used for transforming and assigning event time + to input messages from a source. + """ + + def __call__(self, *args, **kwargs): + """ + This allows to execute the handler function directly if + class instance is sent as a callable. + """ + return self.handler(*args, **kwargs) + + @abstractmethod + async def handler(self, keys: list[str], datum: Datum) -> Messages: + """ + 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. + + Returns: + Messages: A collection of transformed messages with potentially modified + event times and tags for conditional forwarding. + """ + pass + diff --git a/packages/pynumaflow-lite/pynumaflow_lite/sourcetransformer.pyi b/packages/pynumaflow-lite/pynumaflow_lite/sourcetransformer.pyi new file mode 100644 index 00000000..9d7d368c --- /dev/null +++ b/packages/pynumaflow-lite/pynumaflow_lite/sourcetransformer.pyi @@ -0,0 +1,70 @@ +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 ._sourcetransformer_dtypes import SourceTransformer as SourceTransformer + + +class Messages: + def __init__(self) -> None: ... + + def append(self, message: Message) -> None: ... + + def __repr__(self) -> str: ... + + def __str__(self) -> str: ... + + +class Message: + keys: Optional[List[str]] + value: bytes + event_time: _dt.datetime + tags: Optional[List[str]] + + def __init__( + self, + value: bytes, + event_time: _dt.datetime, + keys: Optional[List[str]] = ..., + tags: Optional[List[str]] = ..., + ) -> None: ... + + @staticmethod + def message_to_drop(event_time: _dt.datetime) -> Message: ... + + +class Datum: + # Read-only attributes provided by the extension + keys: List[str] + value: bytes + watermark: _dt.datetime + event_time: _dt.datetime + headers: Dict[str, str] + + def __repr__(self) -> str: ... + + def __str__(self) -> str: ... + + +class SourceTransformAsyncServer: + 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__ = [ + "Messages", + "Message", + "Datum", + "SourceTransformAsyncServer", + "SourceTransformer", +] + diff --git a/packages/pynumaflow-lite/src/lib.rs b/packages/pynumaflow-lite/src/lib.rs index 6b386d11..b3aa76ec 100644 --- a/packages/pynumaflow-lite/src/lib.rs +++ b/packages/pynumaflow-lite/src/lib.rs @@ -9,6 +9,7 @@ pub mod reducestream; pub mod session_reduce; pub mod sink; pub mod source; +pub mod sourcetransform; use pyo3::prelude::*; @@ -75,6 +76,13 @@ fn sourcer(_py: Python, m: &Bound) -> PyResult<()> { Ok(()) } +/// Submodule: pynumaflow_lite.sourcetransformer +#[pymodule] +fn sourcetransformer(_py: Python, m: &Bound) -> PyResult<()> { + crate::sourcetransform::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<()> { @@ -88,6 +96,7 @@ fn pynumaflow_lite(py: Python, m: &Bound) -> PyResult<()> { m.add_wrapped(pyo3::wrap_pymodule!(accumulator))?; m.add_wrapped(pyo3::wrap_pymodule!(sinker))?; m.add_wrapped(pyo3::wrap_pymodule!(sourcer))?; + m.add_wrapped(pyo3::wrap_pymodule!(sourcetransformer))?; // Ensure it's importable as `pynumaflow_lite.mapper` as well as attribute access let binding = m.getattr("mapper")?; @@ -170,5 +179,14 @@ fn pynumaflow_lite(py: Python, m: &Bound) -> PyResult<()> { .getattr("modules")? .set_item(fullname, &sub)?; + // Ensure it's importable as `pynumaflow_lite.sourcetransformer` as well + let binding = m.getattr("sourcetransformer")?; + let sub = binding.cast::()?; + let fullname = "pynumaflow_lite.sourcetransformer"; + sub.setattr("__name__", fullname)?; + py.import("sys")? + .getattr("modules")? + .set_item(fullname, &sub)?; + Ok(()) } diff --git a/packages/pynumaflow-lite/src/sourcetransform/mod.rs b/packages/pynumaflow-lite/src/sourcetransform/mod.rs new file mode 100644 index 00000000..9730ee8f --- /dev/null +++ b/packages/pynumaflow-lite/src/sourcetransform/mod.rs @@ -0,0 +1,228 @@ +use std::collections::HashMap; + +use numaflow::sourcetransform; + +use chrono::{DateTime, Utc}; + +/// SourceTransform 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 collection of [Message]s. +#[pyclass(module = "pynumaflow_lite.sourcetransformer")] +#[derive(Clone, Debug)] +pub struct Messages { + pub(crate) messages: Vec, +} + +#[pymethods] +impl Messages { + #[new] + #[pyo3(signature = () -> "Messages")] + fn new() -> Self { + Self { messages: vec![] } + } + + /// Append a [Message] to the collection. + #[pyo3(signature = (message: "Message"))] + fn append(&mut self, message: Message) { + self.messages.push(message); + } + + fn __repr__(&self) -> String { + format!("Messages({:?})", self.messages) + } + + fn __str__(&self) -> String { + format!("Messages({:?})", self.messages) + } +} + +/// A message to be sent to the next vertex with event time transformation. +#[pyclass(module = "pynumaflow_lite.sourcetransformer")] +#[derive(Clone, Default, Debug)] +pub struct Message { + /// Keys are a collection of strings which will be passed on to the next vertex as is. It can + /// be an empty collection. + pub keys: Option>, + /// Value is the value passed to the next vertex. + pub value: Vec, + /// Time for the given event. This will be used for tracking watermarks. + pub event_time: DateTime, + /// Tags are used for conditional forwarding. + pub tags: Option>, +} + +#[pymethods] +impl Message { + /// Create a new [Message] with the given value, event_time, keys, and tags. + #[new] + #[pyo3(signature = (value: "bytes", event_time: "datetime.datetime", keys: "list[str] | None"=None, tags: "list[str] | None"=None) -> "Message" + )] + fn new( + value: Vec, + event_time: DateTime, + keys: Option>, + tags: Option>, + ) -> Self { + Self { + keys, + value, + event_time, + tags, + } + } + + /// Drop a [Message], do not forward to the next vertex. + /// Event time is required because even though a message is dropped, + /// it is still considered as being processed, hence the watermark should be updated. + #[pyo3(signature = (event_time: "datetime.datetime"))] + #[staticmethod] + fn message_to_drop(event_time: DateTime) -> Self { + Self { + keys: None, + value: vec![], + event_time, + tags: Some(vec![numaflow::shared::DROP.to_string()]), + } + } +} + +impl From for sourcetransform::Message { + fn from(value: Message) -> Self { + Self::new(value.value, value.event_time) + .with_keys(value.keys.unwrap_or_default()) + .with_tags(value.tags.unwrap_or_default()) + } +} + +/// The incoming [SourceTransformRequest] accessible in Python function. +#[pyclass(module = "pynumaflow_lite.sourcetransformer")] +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 event_time: DateTime, + /// Headers for the message. + #[pyo3(get)] + pub headers: HashMap, +} + +impl Datum { + fn new( + keys: Vec, + value: Vec, + watermark: DateTime, + event_time: DateTime, + headers: HashMap, + ) -> Self { + Self { + keys, + value, + watermark, + event_time, + headers, + } + } + + fn __repr__(&self) -> String { + format!( + "Datum(keys={:?}, value={:?}, watermark={}, event_time={}, headers={:?})", + self.keys, self.value, self.watermark, self.event_time, self.headers + ) + } + + fn __str__(&self) -> String { + format!( + "Datum(keys={:?}, value={:?}, watermark={}, event_time={}, headers={:?})", + self.keys, + String::from_utf8_lossy(&self.value), + self.watermark, + self.event_time, + self.headers + ) + } +} + +impl From for Datum { + fn from(value: sourcetransform::SourceTransformRequest) -> Self { + Datum::new( + value.keys, + value.value, + value.watermark, + value.eventtime, + value.headers, + ) + } +} + +/// Async SourceTransform Server that can be started from Python code which will run the Python UDF function. +#[pyclass(module = "pynumaflow_lite.sourcetransformer")] +pub struct SourceTransformAsyncServer { + sock_file: String, + info_file: String, + shutdown_tx: Mutex>>, +} + +#[pymethods] +impl SourceTransformAsyncServer { + #[new] + #[pyo3(signature = (sock_file: "str | None"=sourcetransform::SOCK_ADDR.to_string(), info_file: "str | None"=sourcetransform::SERVER_INFO_FILE.to_string()) -> "SourceTransformAsyncServer" + )] + 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::sourcetransform::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 sourcetransform types/functions. +pub(crate) fn populate_py_module(m: &Bound) -> PyResult<()> { + m.add_class::()?; + m.add_class::()?; + m.add_class::()?; + m.add_class::()?; + + Ok(()) +} diff --git a/packages/pynumaflow-lite/src/sourcetransform/server.rs b/packages/pynumaflow-lite/src/sourcetransform/server.rs new file mode 100644 index 00000000..c554074b --- /dev/null +++ b/packages/pynumaflow-lite/src/sourcetransform/server.rs @@ -0,0 +1,90 @@ +use crate::sourcetransform::{Datum, Messages}; +use numaflow::shared::ServerExtras; +use numaflow::sourcetransform; + +use pyo3::prelude::*; +use std::sync::Arc; + +pub(crate) struct PySourceTransformRunner { + pub(crate) event_loop: Arc>, + pub(crate) py_func: Arc>, +} + +#[tonic::async_trait] +impl sourcetransform::SourceTransformer for PySourceTransformRunner { + async fn transform( + &self, + input: sourcetransform::SourceTransformRequest, + ) -> Vec { + let fut = Python::attach(|py| { + let keys = input.keys.clone(); + let input: Datum = input.into(); + let py_func = self.py_func.clone(); + + let locals = pyo3_async_runtimes::TaskLocals::new(self.event_loop.bind(py).clone()); + + let coro = py_func.call1(py, (keys, input)).unwrap().into_bound(py); + + pyo3_async_runtimes::into_future_with_locals(&locals, coro).unwrap() + }); + + let result = fut.await.unwrap(); + + let result = Python::attach(|py| { + let x: Messages = result.extract(py).unwrap(); + x + }); + + println!("{:?}", result); + + result.messages.into_iter().map(|m| m.into()).collect() + } +} + +// Start the sourcetransform 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_sourcetransform_runner = PySourceTransformRunner { + py_func: Arc::new(py_func), + event_loop: event_loop.clone(), + }; + + let server = numaflow::sourcetransform::Server::new(py_sourcetransform_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 SourceTransform 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/sourcetransform.rs b/packages/pynumaflow-lite/tests/bin/sourcetransform.rs new file mode 100644 index 00000000..46473945 --- /dev/null +++ b/packages/pynumaflow-lite/tests/bin/sourcetransform.rs @@ -0,0 +1,144 @@ +use std::env; +use std::path::PathBuf; + +use numaflow::proto; +use numaflow::proto::source_transformer::source_transform_client::SourceTransformClient; +use tokio::net::UnixStream; +use tokio::sync::mpsc; +use tokio_stream::wrappers::ReceiverStream; +use tonic::transport::Uri; +use tower::service_fn; + +// Simple Rust client binary that exercises the SourceTransform server over Unix Domain Socket. +#[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_SOURCETRANSFORM_SOCK").ok()) + .unwrap_or_else(|| "/tmp/var/run/numaflow/sourcetransform.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 = SourceTransformClient::new(channel); + + let (tx, rx) = mpsc::channel(8); + + // Handshake + let handshake_request = proto::source_transformer::SourceTransformRequest { + request: None, + handshake: Some(proto::source_transformer::Handshake { sot: true }), + }; + tx.send(handshake_request).await.unwrap(); + + let resp = client + .source_transform_fn(ReceiverStream::new(rx)) + .await + .unwrap(); + let mut resp = resp.into_inner(); + + let handshake_response = resp.message().await.unwrap(); + assert!(handshake_response.is_some()); + let handshake_response = handshake_response.unwrap(); + assert!(handshake_response.handshake.is_some()); + + // Request 1 - normal message (event time in 2023) + let request_1 = proto::source_transformer::SourceTransformRequest { + request: Some( + proto::source_transformer::source_transform_request::Request { + id: "1".to_string(), + keys: vec!["first".into(), "second".into()], + value: "hello".into(), + watermark: Some(prost_types::Timestamp::default()), + event_time: Some(prost_types::Timestamp { + seconds: 1672531200, // 2023-01-01 00:00:00 UTC + nanos: 0, + }), + headers: Default::default(), + metadata: None, + }, + ), + handshake: None, + }; + tx.send(request_1).await.unwrap(); + + let actual_response = resp.message().await.unwrap(); + assert!(actual_response.is_some()); + let r = actual_response.unwrap(); + assert_eq!(r.id, "1"); + let msg = &r.results[0]; + assert_eq!(msg.keys.first(), Some(&"first".to_owned())); + assert_eq!(msg.value, "hello".as_bytes()); + // Verify event_time is set + assert!(msg.event_time.is_some()); + + // Request 2 - message to be dropped (event time in 2021) + let request_2 = proto::source_transformer::SourceTransformRequest { + request: Some( + proto::source_transformer::source_transform_request::Request { + id: "2".to_string(), + keys: vec!["third".into(), "fourth".into()], + value: "old_message".into(), + watermark: Some(prost_types::Timestamp::default()), + event_time: Some(prost_types::Timestamp { + seconds: 1609459200, // 2021-01-01 00:00:00 UTC + nanos: 0, + }), + headers: Default::default(), + metadata: None, + }, + ), + handshake: None, + }; + tx.send(request_2).await.unwrap(); + + let actual_response = resp.message().await.unwrap(); + assert!(actual_response.is_some()); + let r = actual_response.unwrap(); + assert_eq!(r.id, "2"); + let msg = &r.results[0]; + assert_eq!(msg.tags, vec![numaflow::shared::DROP.to_string()]); + + // Request 3 - message within 2022 + let request_3 = proto::source_transformer::SourceTransformRequest { + request: Some( + proto::source_transformer::source_transform_request::Request { + id: "3".to_string(), + keys: vec!["fifth".into()], + value: "year_2022_message".into(), + watermark: Some(prost_types::Timestamp::default()), + event_time: Some(prost_types::Timestamp { + seconds: 1656633600, // 2022-07-01 00:00:00 UTC + nanos: 0, + }), + headers: Default::default(), + metadata: None, + }, + ), + handshake: None, + }; + tx.send(request_3).await.unwrap(); + + let actual_response = resp.message().await.unwrap(); + assert!(actual_response.is_some()); + let r = actual_response.unwrap(); + assert_eq!(r.id, "3"); + let msg = &r.results[0]; + assert_eq!(msg.value, "year_2022_message".as_bytes()); + assert!(msg.tags.contains(&"within_year_2022".to_string())); + + // close request stream + drop(tx); + + Ok(()) +} diff --git a/packages/pynumaflow-lite/tests/examples/sourcetransform_event_filter.py b/packages/pynumaflow-lite/tests/examples/sourcetransform_event_filter.py new file mode 100644 index 00000000..e39366f0 --- /dev/null +++ b/packages/pynumaflow-lite/tests/examples/sourcetransform_event_filter.py @@ -0,0 +1,95 @@ +import asyncio +import signal +from datetime import datetime, timezone +from pynumaflow_lite import 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): + """ + A source transformer that filters and routes messages based on event time. + + - Messages before 2022 are dropped + - Messages within 2022 are tagged with "within_year_2022" + - Messages after 2022 are tagged with "after_year_2022" + """ + + async def handler( + self, keys: list[str], datum: sourcetransformer.Datum + ) -> sourcetransformer.Messages: + val = datum.value + 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 + + +# 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): + sock_file = "/tmp/var/run/numaflow/sourcetransform.sock" + server_info_file = "/tmp/var/run/numaflow/sourcetransformer-server-info" + server = sourcetransformer.SourceTransformAsyncServer(sock_file, server_info_file) + + # 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 + try: + server.stop() + except Exception: + pass + return + + +if __name__ == "__main__": + async_handler = EventFilter() + asyncio.run(start(async_handler)) + diff --git a/packages/pynumaflow-lite/tests/test_sourcetransform.py b/packages/pynumaflow-lite/tests/test_sourcetransform.py new file mode 100644 index 00000000..a5182048 --- /dev/null +++ b/packages/pynumaflow-lite/tests/test_sourcetransform.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/sourcetransform.sock") +SERVER_INFO = Path("/tmp/var/run/numaflow/sourcetransformer-server-info") + +SCRIPTS = [ + "sourcetransform_event_filter.py", +] + + +@pytest.mark.parametrize("script", SCRIPTS) +def test_python_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_sourcetransform", + ) +