diff --git a/packages/pynumaflow-lite/Cargo.toml b/packages/pynumaflow-lite/Cargo.toml index 2003c5b3..bf48626a 100644 --- a/packages/pynumaflow-lite/Cargo.toml +++ b/packages/pynumaflow-lite/Cargo.toml @@ -48,6 +48,10 @@ path = "tests/bin/session_reduce.rs" name = "test_accumulator" path = "tests/bin/accumulator.rs" +[[bin]] +name = "test_accumulator_blackhole" +path = "tests/bin/accumulator_blackhole.rs" + [[bin]] name = "test_reducestream" path = "tests/bin/reducestream.rs" diff --git a/packages/pynumaflow-lite/pynumaflow_lite/_accumulator_dtypes.py b/packages/pynumaflow-lite/pynumaflow_lite/_accumulator_dtypes.py index 9178e523..fa3958f5 100644 --- a/packages/pynumaflow-lite/pynumaflow_lite/_accumulator_dtypes.py +++ b/packages/pynumaflow-lite/pynumaflow_lite/_accumulator_dtypes.py @@ -22,7 +22,7 @@ async def handler(self, datums: AsyncIterator[Datum]) -> AsyncIterator[Message]: """ Accumulate can read unordered from the input stream and emit the ordered data to the output stream. Once the watermark (WM) of the output stream progresses, the data in WAL until that WM will be garbage collected. - NOTE: A message can be silently dropped if need be, and it will be cleared from the WAL when the WM progresses. + NOTE: To drop a datum, yield `Message.message_to_drop(datum)` so the WM can progress. Args: datums: An async iterator of Datum objects diff --git a/packages/pynumaflow-lite/pynumaflow_lite/accumulator.pyi b/packages/pynumaflow-lite/pynumaflow_lite/accumulator.pyi index da56a1c7..574dc1f9 100644 --- a/packages/pynumaflow-lite/pynumaflow_lite/accumulator.pyi +++ b/packages/pynumaflow-lite/pynumaflow_lite/accumulator.pyi @@ -25,9 +25,14 @@ class Message: watermark: datetime | None = None, ) -> None: ... @staticmethod - def message_to_drop() -> Message: + def message_to_drop(datum: Datum) -> Message: """ - Drop a Message, do not forward to the next vertex. + Builds a Message from the given Datum with drop tags set, so the message is not + forwarded to the next vertex but still allows the accumulator to advance the watermark + and release tracked state. + + Args: + datum: The input Datum to drop the results for """ ... diff --git a/packages/pynumaflow-lite/src/accumulate/mod.rs b/packages/pynumaflow-lite/src/accumulate/mod.rs index 3cba3ffc..b469627f 100644 --- a/packages/pynumaflow-lite/src/accumulate/mod.rs +++ b/packages/pynumaflow-lite/src/accumulate/mod.rs @@ -68,18 +68,20 @@ impl Message { } } - /// Drop a Message, do not forward to the next vertex. - #[pyo3(signature = ())] + /// Builds a Message from the given Datum with drop tags set, so the message is not + /// forwarded to the next vertex but still allows the accumulator to advance the watermark + /// and release tracked state. + #[pyo3(signature = (datum))] #[staticmethod] - fn message_to_drop() -> Self { + fn message_to_drop(datum: &Datum) -> Self { Self { - keys: None, + keys: Some(datum.keys.clone()), value: vec![], tags: Some(vec![numaflow::shared::DROP.to_string()]), - id: String::new(), - headers: HashMap::new(), - event_time: chrono::Utc::now(), - watermark: chrono::Utc::now(), + id: datum.id.clone(), + headers: datum.headers.clone(), + event_time: datum.event_time, + watermark: datum.watermark, } } diff --git a/packages/pynumaflow-lite/tests/bin/accumulator_blackhole.rs b/packages/pynumaflow-lite/tests/bin/accumulator_blackhole.rs new file mode 100644 index 00000000..1a8bac70 --- /dev/null +++ b/packages/pynumaflow-lite/tests/bin/accumulator_blackhole.rs @@ -0,0 +1,135 @@ +use std::collections::HashMap; +use std::env; +use std::path::PathBuf; +use std::time::{SystemTime, UNIX_EPOCH}; + +use tokio::net::UnixStream; +use tokio::sync::mpsc; +use tokio_stream::wrappers::ReceiverStream; +use tonic::{Request, transport::Uri}; +use tower::service_fn; + +use numaflow::proto::accumulator as acc_proto; +use numaflow::proto::accumulator::accumulator_request::window_operation::Event; + +fn ts_from_secs(secs: i64) -> prost_types::Timestamp { + prost_types::Timestamp { + seconds: secs, + nanos: 0, + } +} + +fn keyed_window(base_time: i64) -> acc_proto::KeyedWindow { + acc_proto::KeyedWindow { + start: Some(ts_from_secs(base_time)), + end: Some(ts_from_secs(base_time + 60)), + slot: "slot-0".to_string(), + keys: vec!["key1".into()], + } +} + +#[tokio::main] +async fn main() -> Result<(), Box> { + let sock_path = env::args() + .nth(1) + .or_else(|| env::var("NUMAFLOW_ACCUMULATOR_SOCK").ok()) + .unwrap_or_else(|| "/tmp/var/run/numaflow/accumulator.sock".to_string()); + + 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 = + numaflow::proto::accumulator::accumulator_client::AccumulatorClient::new(channel); + + let (tx, rx) = mpsc::channel(16); + + let base_time = SystemTime::now().duration_since(UNIX_EPOCH)?.as_secs() as i64; + + // (id, event_time offset, watermark offset) + let inputs = [("msg1", 30, 5), ("msg2", 10, 15), ("msg3", 20, 25)]; + let headers = HashMap::from([("h1".to_string(), "v1".to_string())]); + + for (i, (id, et, wm)) in inputs.iter().enumerate() { + let event = if i == 0 { Event::Open } else { Event::Append }; + tx.send(acc_proto::AccumulatorRequest { + payload: Some(acc_proto::Payload { + keys: vec!["key1".into()], + value: format!("value-{id}").into_bytes(), + watermark: Some(ts_from_secs(base_time + wm)), + event_time: Some(ts_from_secs(base_time + et)), + headers: headers.clone(), + id: id.to_string(), + }), + operation: Some(acc_proto::accumulator_request::WindowOperation { + event: event as i32, + keyed_window: Some(keyed_window(base_time)), + }), + }) + .await?; + } + + tx.send(acc_proto::AccumulatorRequest { + payload: None, + operation: Some(acc_proto::accumulator_request::WindowOperation { + event: Event::Close as i32, + keyed_window: Some(keyed_window(base_time)), + }), + }) + .await?; + drop(tx); + + let request = Request::new(ReceiverStream::new(rx)); + let mut resp = client.accumulate_fn(request).await?.into_inner(); + + let mut dropped = Vec::new(); + let mut found_eof = false; + + while let Some(r) = resp.message().await? { + if r.eof { + found_eof = true; + continue; + } + assert_eq!( + r.tags, + vec![numaflow::shared::DROP.to_string()], + "Every response should carry the DROP tag" + ); + let payload = r.payload.expect("Drop response should carry a payload"); + assert!( + payload.value.is_empty(), + "Drop message value should be empty" + ); + assert_eq!(payload.keys, vec!["key1".to_string()]); + assert_eq!(payload.headers, headers); + dropped.push(payload); + } + + assert!(found_eof, "Should have received EOF"); + assert_eq!(dropped.len(), inputs.len(), "Expected one drop per datum"); + + for (payload, (id, et, wm)) in dropped.iter().zip(inputs.iter()) { + assert_eq!(payload.id, *id, "Drop message id should match the datum"); + assert_eq!( + payload.event_time, + Some(ts_from_secs(base_time + et)), + "Drop message event_time should match the datum" + ); + assert_eq!( + payload.watermark, + Some(ts_from_secs(base_time + wm)), + "Drop message watermark should match the datum" + ); + } + + println!("All datums were dropped with their metadata preserved!"); + + Ok(()) +} diff --git a/packages/pynumaflow-lite/tests/examples/accumulator_blackhole.py b/packages/pynumaflow-lite/tests/examples/accumulator_blackhole.py new file mode 100644 index 00000000..06be6497 --- /dev/null +++ b/packages/pynumaflow-lite/tests/examples/accumulator_blackhole.py @@ -0,0 +1,49 @@ +""" +Blackhole accumulator example. + +This accumulator discards every datum it receives. Instead of yielding nothing, it yields a +drop message for each datum so the watermark can progress and the tracked state can be released. +""" + +import asyncio +import signal +from collections.abc import AsyncIterator + +from pynumaflow_lite.accumulator import ( + Accumulator, + AccumulatorAsyncServer, + Datum, + Message, +) + + +class Blackhole(Accumulator): + async def handler(self, datums: AsyncIterator[Datum]) -> AsyncIterator[Message]: + async for datum in datums: + print(f"Dropping datum: id={datum.id}, event_time={datum.event_time}") + yield Message.message_to_drop(datum) + + +async def main(): + sock_file = "/tmp/var/run/numaflow/accumulator.sock" + server_info_file = "/tmp/var/run/numaflow/accumulator-server-info" + server = AccumulatorAsyncServer(sock_file, server_info_file) + + 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: + print("Starting Blackhole Accumulator Server...") + await server.start(Blackhole) + print("Shutting down gracefully...") + except asyncio.CancelledError: + server.stop() + return + + +if __name__ == "__main__": + asyncio.run(main()) diff --git a/packages/pynumaflow-lite/tests/test_accumulator.py b/packages/pynumaflow-lite/tests/test_accumulator.py index eeaa9a5b..f868e0d0 100644 --- a/packages/pynumaflow-lite/tests/test_accumulator.py +++ b/packages/pynumaflow-lite/tests/test_accumulator.py @@ -21,3 +21,13 @@ def test_python_accumulator_server_and_rust_client(script: str, tmp_path: Path): rust_bin_name="test_accumulator", rust_bin_args=[str(SOCK_PATH)], ) + + +def test_python_accumulator_drop_messages(tmp_path: Path): + run_python_server_with_rust_client( + script="accumulator_blackhole.py", + sock_path=SOCK_PATH, + server_info_path=SERVER_INFO, + rust_bin_name="test_accumulator_blackhole", + rust_bin_args=[str(SOCK_PATH)], + )