diff --git a/packages/pynumaflow-lite/.gitignore b/packages/pynumaflow-lite/.gitignore new file mode 100644 index 00000000..77d1d50a --- /dev/null +++ b/packages/pynumaflow-lite/.gitignore @@ -0,0 +1 @@ +.docker-cache/ diff --git a/packages/pynumaflow-lite/README.md b/packages/pynumaflow-lite/README.md index 0af3fcd7..f5db0535 100644 --- a/packages/pynumaflow-lite/README.md +++ b/packages/pynumaflow-lite/README.md @@ -24,7 +24,12 @@ make test Go to `pynumaflow-lite` (top level) directory and run the below command. ```bash -docker run --rm -v $(pwd):/io ghcr.io/pyo3/maturin build -i python3.11 --release +docker run --rm \ + -v $(pwd):/io \ + -v $(pwd)/.docker-cache/.cargo-docker-registry:/root/.cargo/registry \ + -v $(pwd)/.docker-cache/cargo-docker-git:/root/.cargo/git \ + -v $(pwd)/target:/io/target \ + ghcr.io/pyo3/maturin build -i python3.11 --release ``` This will create the `wheel` file in `target/wheels/` directory. You should copy it over to where we @@ -33,5 +38,5 @@ are writing the python code referencing this library. e.g., ```bash -cp target/wheels/pynumaflow_lite-0.1.0-cp311-cp311-manylinux_2_17_aarch64.manylinux2014_aarch64.whl manifests/simple-async-map/ -``` \ No newline at end of file +cp -v target/wheels/pynumaflow_lite-0.1.0a1-cp311-cp311-manylinux_2_17_aarch64.manylinux2014_aarch64.whl manifests/map/ +``` diff --git a/packages/pynumaflow-lite/manifests/map/.dockerignore b/packages/pynumaflow-lite/manifests/map/.dockerignore new file mode 100644 index 00000000..21d0b898 --- /dev/null +++ b/packages/pynumaflow-lite/manifests/map/.dockerignore @@ -0,0 +1 @@ +.venv/ diff --git a/packages/pynumaflow-lite/manifests/map/Dockerfile b/packages/pynumaflow-lite/manifests/map/Dockerfile index ef530774..76aa4bda 100644 --- a/packages/pynumaflow-lite/manifests/map/Dockerfile +++ b/packages/pynumaflow-lite/manifests/map/Dockerfile @@ -1,37 +1,42 @@ -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", "map_cat.py"] \ No newline at end of file +FROM python:3.11-slim-trixie AS builder + +COPY --from=ghcr.io/astral-sh/uv:latest /uv /uvx /bin/ + +ENV UV_COMPILE_BYTECODE=1 \ + UV_LINK_MODE=copy \ + UV_PYTHON_DOWNLOADS=never + +WORKDIR /app + +RUN --mount=type=cache,target=/root/.cache/uv \ + --mount=type=bind,source=uv.lock,target=uv.lock \ + --mount=type=bind,source=pyproject.toml,target=pyproject.toml \ + uv sync --locked --no-install-project --no-dev + +COPY . . + +RUN --mount=type=cache,target=/root/.cache/uv \ + uv sync --locked --no-dev + +RUN uv pip install ./pynumaflow_lite-0.1.0a1-cp311-cp311-manylinux_2_17_aarch64.manylinux2014_aarch64.whl + +FROM python:3.11-slim-trixie + +# Setup a non-root user +RUN groupadd --system --gid 999 nonroot \ + && useradd --system --gid 999 --uid 999 --create-home nonroot + +COPY --from=builder --chown=nonroot:nonroot /app /app + +ENV PATH="/app/.venv/bin:$PATH" + +# Keeps Python from buffering stdout and stderr to avoid situations where +# the application crashes without emitting any logs due to buffering. +ENV PYTHONUNBUFFERED=1 + +# Use the non-root user to run our application +USER nonroot + +WORKDIR /app + +CMD ["python", "map_cat.py"] diff --git a/packages/pynumaflow-lite/manifests/map/README.md b/packages/pynumaflow-lite/manifests/map/README.md index ecef703d..6f826d88 100644 --- a/packages/pynumaflow-lite/manifests/map/README.md +++ b/packages/pynumaflow-lite/manifests/map/README.md @@ -16,4 +16,4 @@ k3d image import quay.io/numaio/numaflow/pynumaflow-lite-map-cat:v2 ```bash kubectl apply -f pipeline.yaml -``` \ No newline at end of file +``` diff --git a/packages/pynumaflow-lite/manifests/map/map_cat.py b/packages/pynumaflow-lite/manifests/map/map_cat.py index f877f212..4a4811af 100644 --- a/packages/pynumaflow-lite/manifests/map/map_cat.py +++ b/packages/pynumaflow-lite/manifests/map/map_cat.py @@ -1,48 +1,22 @@ import asyncio -import signal -from collections.abc import Awaitable, Callable -from pynumaflow_lite import mapper +from pynumaflow_lite.mapper import Datum, MapAsyncServer, Mapper, Message -class SimpleCat(mapper.Mapper): - async def handler(self, keys: list[str], payload: mapper.Datum) -> mapper.Messages: +class SimpleCat(Mapper): + async def handler(self, datum: Datum) -> list[Message]: + if datum.value == b"bad world": + return [Message.to_drop()] + print(f"Received {datum=}") + return [Message(datum.value, keys=datum.keys)] - messages = mapper.Messages() - if payload.value == b"bad world": - messages.append(mapper.Message.message_to_drop()) - else: - messages.append(mapper.Message(payload.value, keys)) - - return messages - - -async def start(f: Callable[[list[str], mapper.Datum], Awaitable[mapper.Messages]]): - server = mapper.MapAsyncServer() - - # 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 - server.stop() - return +async def main() -> None: + print("Starting map server") + # `serve` returns when SIGINT or SIGTERM arrives. + await MapAsyncServer(SimpleCat()).serve() + print("Map server stopped") if __name__ == "__main__": - async_handler = SimpleCat() - asyncio.run(start(async_handler)) + asyncio.run(main()) diff --git a/packages/pynumaflow-lite/manifests/map/pipeline.yaml b/packages/pynumaflow-lite/manifests/map/pipeline.yaml index 25311f2e..90f48c2e 100644 --- a/packages/pynumaflow-lite/manifests/map/pipeline.yaml +++ b/packages/pynumaflow-lite/manifests/map/pipeline.yaml @@ -8,8 +8,8 @@ spec: source: # A self data generating source generator: - rpu: 500 - duration: 1s + rpu: 1 + duration: 3s - name: map partitions: 2 scale: @@ -27,4 +27,4 @@ spec: - from: in to: map - from: map - to: sink \ No newline at end of file + to: sink diff --git a/packages/pynumaflow-lite/manifests/map/pyproject.toml b/packages/pynumaflow-lite/manifests/map/pyproject.toml index cc1f0ad2..fd622e46 100644 --- a/packages/pynumaflow-lite/manifests/map/pyproject.toml +++ b/packages/pynumaflow-lite/manifests/map/pyproject.toml @@ -6,11 +6,6 @@ authors = [ { name = "Vigith Maurice", email = "vigith@gmail.com" } ] readme = "README.md" -requires-python = ">=3.11" +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/map/uv.lock b/packages/pynumaflow-lite/manifests/map/uv.lock new file mode 100644 index 00000000..2d5af88c --- /dev/null +++ b/packages/pynumaflow-lite/manifests/map/uv.lock @@ -0,0 +1,8 @@ +version = 1 +revision = 3 +requires-python = "==3.11.*" + +[[package]] +name = "map-cat" +version = "0.1.0" +source = { virtual = "." } diff --git a/packages/pynumaflow-lite/pynumaflow_lite/__init__.py b/packages/pynumaflow-lite/pynumaflow_lite/__init__.py index d1bfeb9f..46a781e2 100644 --- a/packages/pynumaflow-lite/pynumaflow_lite/__init__.py +++ b/packages/pynumaflow-lite/pynumaflow_lite/__init__.py @@ -66,6 +66,7 @@ from ._accumulator_dtypes import Accumulator from ._batchmapper_dtypes import BatchMapper from ._map_dtypes import Mapper +from ._map_server import MapAsyncServer from ._mapstream_dtypes import MapStreamer from ._reduce_dtypes import Reducer from ._reducestreamer_dtypes import ReduceStreamer @@ -78,6 +79,7 @@ if mapper is not None: mapper.Mapper = Mapper + mapper.MapAsyncServer = MapAsyncServer if batchmapper is not None: batchmapper.BatchMapper = BatchMapper diff --git a/packages/pynumaflow-lite/pynumaflow_lite/_map_dtypes.py b/packages/pynumaflow-lite/pynumaflow_lite/_map_dtypes.py index 13733ed7..cb92c551 100644 --- a/packages/pynumaflow-lite/pynumaflow_lite/_map_dtypes.py +++ b/packages/pynumaflow-lite/pynumaflow_lite/_map_dtypes.py @@ -1,6 +1,6 @@ from abc import ABCMeta, abstractmethod -from pynumaflow_lite.mapper import Datum, Messages +from pynumaflow_lite.mapper import Datum, Message class Mapper(metaclass=ABCMeta): @@ -17,7 +17,7 @@ class instance is sent as a callable. return self.handler(*args, **kwargs) @abstractmethod - async def handler(self, keys: list[str], payload: Datum) -> Messages: + async def handler(self, payload: Datum) -> list[Message]: """ Implement this handler function which implements the MapAsyncCallable interface. """ diff --git a/packages/pynumaflow-lite/pynumaflow_lite/_map_server.py b/packages/pynumaflow-lite/pynumaflow_lite/_map_server.py new file mode 100644 index 00000000..c758051b --- /dev/null +++ b/packages/pynumaflow-lite/pynumaflow_lite/_map_server.py @@ -0,0 +1,131 @@ +from __future__ import annotations + +import asyncio +import contextlib +import signal +from collections.abc import Awaitable, Callable +from types import TracebackType + +from .pynumaflow_lite import mapper as _mapper + +Datum = _mapper.Datum +Message = _mapper.Message + +_SHUTDOWN_SIGNALS = (signal.SIGINT, signal.SIGTERM) + + +class MapAsyncServer: + def __init__( + self, + handler: Callable[[Datum], Awaitable[list[Message]]], + *, + sock_file: str | None = None, + server_info_file: str | None = None, + install_signal_handlers: bool = True, + ) -> None: + self._core = _mapper._MapAsyncServer(sock_file, server_info_file) + self._handler = handler + self._install_signal_handlers = install_signal_handlers + self._task: asyncio.Task[None] | None = None + self._serving = False + self._installed_signals: list[signal.Signals] = [] + + async def serve(self) -> None: + """Run the map server until it stops. + + This is the entrypoint for an application that already runs an event + loop. It returns when a shutdown signal arrives or when `stop()` runs. + """ + await self._serve(install_signal_handlers=self._install_signal_handlers) + + async def _serve(self, *, install_signal_handlers: bool) -> None: + if self._serving: + raise RuntimeError("map server is already serving") + self._serving = True + try: + if install_signal_handlers: + self._add_signal_handlers() + await self._core.start(self._handler) + finally: + self._remove_signal_handlers() + self._serving = False + + def stop(self) -> None: + self._core.stop() + + async def wait_ready(self, timeout: float = 30.0) -> None: + await self._core.wait_ready(timeout) + + async def wait_for_termination(self) -> None: + """Wait until the background server task ends. + + Use this inside an `async with` block. It raises the handler error if + the server task failed. + """ + if self._task is None: + raise RuntimeError("map server is not serving") + await asyncio.shield(self._task) + + def _add_signal_handlers(self) -> None: + loop = asyncio.get_running_loop() + for sig in _SHUTDOWN_SIGNALS: + try: + loop.add_signal_handler(sig, self.stop) + except (NotImplementedError, RuntimeError, OSError): + continue + self._installed_signals.append(sig) + + def _remove_signal_handlers(self) -> None: + if not self._installed_signals: + return + try: + loop = asyncio.get_running_loop() + except RuntimeError: + self._installed_signals.clear() + return + for sig in self._installed_signals: + with contextlib.suppress(NotImplementedError, OSError): + loop.remove_signal_handler(sig) + self._installed_signals.clear() + + async def __aenter__(self) -> MapAsyncServer: + """Start the server in a background task and wait until it is ready. + + This form is for tests and for code that must run other work next to + the server. It never installs signal handlers. + """ + if self._task is not None and not self._task.done(): + raise RuntimeError("map server is already serving") + + self._task = asyncio.create_task(self._serve(install_signal_handlers=False)) + try: + await self.wait_ready() + except BaseException: + self.stop() + task, self._task = self._task, None + if task is not None: + # Surface the server error, if there is one. It explains the + # failure better than the `wait_ready` error does. + await task + raise + return self + + async def __aexit__( + self, + exc_type: type[BaseException] | None, + exc: BaseException | None, + tb: TracebackType | None, + ) -> None: + self.stop() + if self._task is not None: + try: + await self._task + finally: + self._task = None + + def run(self) -> None: + """Run the map server in a new event loop until it stops.""" + try: + asyncio.run(self.serve()) + except KeyboardInterrupt: + self.stop() diff --git a/packages/pynumaflow-lite/pynumaflow_lite/mapper.pyi b/packages/pynumaflow-lite/pynumaflow_lite/mapper.pyi index 8255e6d6..c82670b5 100644 --- a/packages/pynumaflow-lite/pynumaflow_lite/mapper.pyi +++ b/packages/pynumaflow-lite/pynumaflow_lite/mapper.pyi @@ -2,114 +2,89 @@ from __future__ import annotations import datetime as _dt from collections.abc import Awaitable, Callable +from types import TracebackType -# Re-export the Python ABC for user convenience and typing from ._map_dtypes import Mapper as Mapper -class SystemMetadata: - """System-generated metadata groups per message (read-only).""" - - def __init__(self) -> None: ... - def groups(self) -> list[str]: - """Returns the groups of the system metadata.""" - ... - - def keys(self, group: str) -> list[str]: - """Returns the keys of the system metadata for the given group.""" - ... - - def value(self, group: str, key: str) -> bytes: - """Returns the value of the system metadata for the given group and key.""" - ... - - def __repr__(self) -> str: ... - -class UserMetadata: - """User-defined metadata groups per message (read-write).""" - - def __init__(self) -> None: ... - def groups(self) -> list[str]: - """Returns the groups of the user metadata.""" - ... - - def keys(self, group: str) -> list[str]: - """Returns the keys of the user metadata for the given group.""" - ... - - def value(self, group: str, key: str) -> bytes: - """Returns the value of the user metadata for the given group and key.""" - ... - - def create_group(self, group: str) -> None: - """Creates a new group in the user metadata.""" - ... - - def add_kv(self, group: str, key: str, value: bytes) -> None: - """Adds a key-value pair to the user metadata.""" - ... - - def remove_key(self, group: str, key: str) -> None: - """Removes a key from a group in the user metadata.""" - ... - - def remove_group(self, group: str) -> None: - """Removes a group from the user metadata.""" - ... - - def __repr__(self) -> str: ... - -class Messages: - def __init__(self) -> None: ... - def append(self, message: Message) -> None: ... - def __repr__(self) -> str: ... - def __str__(self) -> str: ... - class Message: keys: list[str] | None value: bytes tags: list[str] | None - user_metadata: UserMetadata | None + user_metadata: dict[str, dict[str, bytes]] | None def __init__( self, value: bytes, keys: list[str] | None = ..., tags: list[str] | None = ..., - user_metadata: UserMetadata | None = ..., + user_metadata: dict[str, dict[str, bytes]] | None = ..., ) -> None: ... @staticmethod - def message_to_drop() -> Message: ... + def to_drop() -> Message: ... + def __repr__(self) -> str: ... + def __eq__(self, other: object) -> bool: ... class Datum: - # Read-only attributes provided by the extension keys: list[str] value: bytes watermark: _dt.datetime - eventtime: _dt.datetime + event_time: _dt.datetime headers: dict[str, str] - user_metadata: UserMetadata - system_metadata: SystemMetadata + user_metadata: dict[str, dict[str, bytes]] + system_metadata: dict[str, dict[str, bytes]] + def __init__( + self, + *, + keys: list[str] = ..., + value: bytes = ..., + event_time: _dt.datetime | None = ..., + watermark: _dt.datetime | None = ..., + headers: dict[str, str] = ..., + user_metadata: dict[str, dict[str, bytes]] = ..., + system_metadata: dict[str, dict[str, bytes]] = ..., + ) -> None: ... def __repr__(self) -> str: ... def __str__(self) -> str: ... -class MapAsyncServer: +class _MapAsyncServer: def __init__( self, sock_file: str | None = ..., - info_file: str | None = ..., + server_info_file: str | None = ..., ) -> None: ... - def start(self, py_func: Callable[[list[str], Datum], Awaitable[Messages]]) -> Awaitable[None]: ... + def start( + self, + handler: Callable[[Datum], Awaitable[list[Message]]], + ) -> Awaitable[None]: ... + def wait_ready(self, timeout: float = ...) -> Awaitable[None]: ... def stop(self) -> None: ... -# Simple utility function exposed by the extension +class MapAsyncServer: + def __init__( + self, + handler: Callable[[Datum], Awaitable[list[Message]]], + *, + sock_file: str | None = ..., + server_info_file: str | None = ..., + install_signal_handlers: bool = ..., + ) -> None: ... + def run(self) -> None: ... + async def serve(self) -> None: ... + def stop(self) -> None: ... + async def wait_ready(self, timeout: float = ...) -> None: ... + async def wait_for_termination(self) -> None: ... + async def __aenter__(self) -> MapAsyncServer: ... + async def __aexit__( + self, + exc_type: type[BaseException] | None, + exc: BaseException | None, + tb: TracebackType | None, + ) -> None: ... __all__ = [ "Datum", "MapAsyncServer", "Mapper", "Message", - "Messages", - "SystemMetadata", - "UserMetadata", ] diff --git a/packages/pynumaflow-lite/src/map/mod.rs b/packages/pynumaflow-lite/src/map/mod.rs index 85c3843c..0232cf00 100644 --- a/packages/pynumaflow-lite/src/map/mod.rs +++ b/packages/pynumaflow-lite/src/map/mod.rs @@ -1,6 +1,9 @@ use std::collections::HashMap; +use std::path::PathBuf; +use std::time::Duration; use numaflow::map; +use numaflow::proto::map::map_client::MapClient; use chrono::{DateTime, Utc}; @@ -8,233 +11,89 @@ use chrono::{DateTime, Utc}; /// and can pass in the Python function. pub mod server; +use tokio::net::UnixStream; +use tokio::time::Instant; +use tonic::transport::Uri; +use tower::service_fn; + use pyo3::prelude::*; use std::sync::Mutex; -/// SystemMetadata wraps system-generated metadata groups per message. -/// It is read-only to UDFs. -#[pyclass(module = "pynumaflow_lite.mapper", from_py_object)] -#[derive(Clone, Default, Debug)] -pub struct SystemMetadata { - data: HashMap>>, -} - -#[pymethods] -impl SystemMetadata { - #[new] - #[pyo3(signature = () -> "SystemMetadata")] - fn new() -> Self { - Self::default() - } - - /// Returns the groups of the system metadata. - /// If there are no groups, it returns an empty list. - #[pyo3(signature = () -> "list[str]")] - fn groups(&self) -> Vec { - self.data.keys().cloned().collect() - } - - /// Returns the keys of the system metadata for the given group. - /// If there are no keys or the group is not present, it returns an empty list. - #[pyo3(signature = (group: "str") -> "list[str]")] - fn keys(&self, group: &str) -> Vec { - self.data - .get(group) - .map(|kv| kv.keys().cloned().collect()) - .unwrap_or_default() - } - - /// Returns the value of the system metadata for the given group and key. - /// If there is no value or the group or key is not present, it returns an empty bytes. - #[pyo3(signature = (group: "str", key: "str") -> "bytes")] - fn value(&self, group: &str, key: &str) -> Vec { - self.data - .get(group) - .and_then(|kv| kv.get(key)) - .cloned() - .unwrap_or_default() - } - - fn __repr__(&self) -> String { - format!("SystemMetadata(groups={:?})", self.groups()) - } -} - -impl From for SystemMetadata { - fn from(value: map::SystemMetadata) -> Self { - let mut data = HashMap::new(); - for group in value.groups() { - let mut kv = HashMap::new(); - for key in value.keys(&group) { - kv.insert(key.clone(), value.value(&group, &key)); - } - data.insert(group, kv); - } - Self { data } - } -} - -/// UserMetadata wraps user-defined metadata groups per message. -/// Users can read and write to this metadata. -#[pyclass(module = "pynumaflow_lite.mapper", from_py_object)] -#[derive(Clone, Default, Debug)] -pub struct UserMetadata { - data: HashMap>>, +pub(crate) fn bytes_literal(value: &[u8]) -> String { + format!("b\"{}\"", String::from_utf8_lossy(value).escape_debug()) } -#[pymethods] -impl UserMetadata { - #[new] - #[pyo3(signature = () -> "UserMetadata")] - fn new() -> Self { - Self::default() - } - - /// Returns the groups of the user metadata. - /// If there are no groups, it returns an empty list. - #[pyo3(signature = () -> "list[str]")] - fn groups(&self) -> Vec { - self.data.keys().cloned().collect() - } - - /// Returns the keys of the user metadata for the given group. - /// If there are no keys or the group is not present, it returns an empty list. - #[pyo3(signature = (group: "str") -> "list[str]")] - fn keys(&self, group: &str) -> Vec { - self.data - .get(group) - .map(|kv| kv.keys().cloned().collect()) - .unwrap_or_default() - } - - /// Returns the value of the user metadata for the given group and key. - /// If there is no value or the group or key is not present, it returns an empty bytes. - #[pyo3(signature = (group: "str", key: "str") -> "bytes")] - fn value(&self, group: &str, key: &str) -> Vec { - self.data - .get(group) - .and_then(|kv| kv.get(key)) - .cloned() - .unwrap_or_default() - } - - /// Creates a new group in the user metadata. - /// If the group already exists, this is a no-op. - #[pyo3(signature = (group: "str"))] - fn create_group(&mut self, group: String) { - self.data.entry(group).or_default(); - } - - /// Adds a key-value pair to the user metadata. - /// If the group is not present, it creates a new group. - #[pyo3(signature = (group: "str", key: "str", value: "bytes"))] - fn add_kv(&mut self, group: String, key: String, value: Vec) { - self.data.entry(group).or_default().insert(key, value); - } - - /// Removes a key from a group in the user metadata. - /// If the key or group is not present, it's a no-op. - #[pyo3(signature = (group: "str", key: "str"))] - fn remove_key(&mut self, group: &str, key: &str) { - if let Some(kv) = self.data.get_mut(group) { - kv.remove(key); - } - } - - /// Removes a group from the user metadata. - /// If the group is not present, it's a no-op. - #[pyo3(signature = (group: "str"))] - fn remove_group(&mut self, group: &str) { - self.data.remove(group); - } - - fn __repr__(&self) -> String { - format!("UserMetadata(groups={:?})", self.groups()) - } +fn metadata_literal(metadata: &HashMap>>) -> String { + let groups: Vec = metadata + .iter() + .map(|(group, kv)| { + let entries: Vec = kv + .iter() + .map(|(key, value)| format!("{:?}: {}", key, bytes_literal(value))) + .collect(); + format!("{:?}: {{{}}}", group, entries.join(", ")) + }) + .collect(); + format!("{{{}}}", groups.join(", ")) } -impl From for UserMetadata { - fn from(value: map::UserMetadata) -> Self { - let mut data = HashMap::new(); - for group in value.groups() { - let mut kv = HashMap::new(); - for key in value.keys(&group) { - kv.insert(key.clone(), value.value(&group, &key)); - } - data.insert(group, kv); +fn system_metadata_to_hash_map( + value: map::SystemMetadata, +) -> HashMap>> { + let mut data = HashMap::new(); + for group in value.groups() { + let mut kv = HashMap::new(); + for key in value.keys(&group) { + kv.insert(key.clone(), value.value(&group, &key)); } - Self { data } + data.insert(group, kv); } + data } -impl From for map::UserMetadata { - fn from(value: UserMetadata) -> Self { - let mut user_metadata = map::UserMetadata::new(); - for (group, kv_map) in value.data { - for (key, val) in kv_map { - user_metadata.add_kv(group.clone(), key, val); - } +fn user_metadata_to_hash_map( + value: map::UserMetadata, +) -> HashMap>> { + let mut data = HashMap::new(); + for group in value.groups() { + let mut kv = HashMap::new(); + for key in value.keys(&group) { + kv.insert(key.clone(), value.value(&group, &key)); } - user_metadata - } -} - -/// A collection of [Message]s. -#[pyclass(module = "pynumaflow_lite.mapper", from_py_object)] -#[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) + data.insert(group, kv); } + data } /// A message to be sent to the next vertex. -#[pyclass(module = "pynumaflow_lite.mapper", from_py_object)] -#[derive(Clone, Default, Debug)] +#[pyclass(module = "pynumaflow_lite.mapper", from_py_object, eq)] +#[derive(Clone, Default, Debug, PartialEq)] 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. + #[pyo3(get)] pub keys: Option>, /// Value is the value passed to the next vertex. + #[pyo3(get)] pub value: Vec, /// Tags are used for [conditional forwarding](https://numaflow.numaproj.io/user-guide/reference/conditional-forwarding/). + #[pyo3(get)] pub tags: Option>, /// User metadata for the message. - pub user_metadata: Option, + #[pyo3(get)] + pub user_metadata: Option>>>, } #[pymethods] impl Message { - /// Create a new [Message] with the given value, keys, tags, and user_metadata. + /// Create a new Message with the given value. Keys, tags, and user_metadata are optional. #[new] - #[pyo3(signature = (value: "bytes", keys: "list[str] | None"=None, tags: "list[str] | None"=None, user_metadata: "UserMetadata | None"=None) -> "Message" - )] + #[pyo3(signature = (value: "bytes", keys: "list[str] | None"=None, tags: "list[str] | None"=None, user_metadata: "dict[str, dict[str, bytes]] | None"=None) -> "Message")] fn new( value: Vec, keys: Option>, tags: Option>, - user_metadata: Option, + user_metadata: Option>>>, ) -> Self { Self { keys, @@ -244,10 +103,10 @@ impl Message { } } - /// Drop a [Message], do not forward to the next vertex. - #[pyo3(signature = ())] + /// A Message marked to be dropped, i.e. not forwarded to the next vertex. #[staticmethod] - fn message_to_drop() -> Self { + #[pyo3(signature = () -> "Message")] + fn to_drop() -> Self { Self { keys: None, value: vec![], @@ -255,6 +114,22 @@ impl Message { user_metadata: None, } } + + fn __repr__(&self) -> String { + format!( + "Message(value={}, keys={}, tags={}, user_metadata={})", + bytes_literal(&self.value), + self.keys + .as_ref() + .map_or_else(|| "None".to_string(), |keys| format!("{keys:?}")), + self.tags + .as_ref() + .map_or_else(|| "None".to_string(), |tags| format!("{tags:?}")), + self.user_metadata + .as_ref() + .map_or_else(|| "None".to_string(), metadata_literal), + ) + } } impl From for map::Message { @@ -263,12 +138,21 @@ impl From for map::Message { keys: value.keys, value: value.value, tags: value.tags, - user_metadata: value.user_metadata.map(|m| m.into()), + user_metadata: value.user_metadata.map(|m| { + let mut umd = map::UserMetadata::new(); + for (group, kv) in m { + for (key, val) in kv { + umd.add_kv(group.clone(), key, val); + } + } + umd + }), } } } -/// The incoming [MapRequest] accessible in Python function. +/// The incoming Datum passed to the map handler. It carries the event's keys, value, +/// event_time, watermark, headers, and the user/system metadata. #[pyclass(module = "pynumaflow_lite.mapper")] pub struct Datum { /// Set of keys in the (key, value) terminology of map/reduce paradigm. @@ -283,106 +167,162 @@ pub struct Datum { pub watermark: DateTime, /// Time of the element as seen at source or aligned after a reduce operation. #[pyo3(get)] - pub eventtime: DateTime, + pub event_time: DateTime, /// Headers for the message. #[pyo3(get)] pub headers: HashMap, /// User metadata for the message. #[pyo3(get)] - pub user_metadata: UserMetadata, + pub user_metadata: HashMap>>, /// System metadata for the message. #[pyo3(get)] - pub system_metadata: SystemMetadata, + pub system_metadata: HashMap>>, } +#[pymethods] impl Datum { + #[new] + #[pyo3(signature = ( + *, + keys: "list[str] | None"=None, + value: "bytes | None"=None, + event_time: "datetime.datetime | None"=None, + watermark: "datetime.datetime | None"=None, + headers: "dict[str, str] | None"=None, + user_metadata: "dict[str, dict[str, bytes]] | None"=None, + system_metadata: "dict[str, dict[str, bytes]] | None"=None, + ) -> "Datum")] fn new( - keys: Vec, - value: Vec, - watermark: DateTime, - eventtime: DateTime, - headers: HashMap, - user_metadata: UserMetadata, - system_metadata: SystemMetadata, + keys: Option>, + value: Option>, + event_time: Option>, + watermark: Option>, + headers: Option>, + user_metadata: Option>>>, + system_metadata: Option>>>, ) -> Self { Self { - keys, - value, - watermark, - eventtime, - headers, - user_metadata, - system_metadata, + keys: keys.unwrap_or_default(), + value: value.unwrap_or_default(), + watermark: watermark.unwrap_or(DateTime::::UNIX_EPOCH), + event_time: event_time.unwrap_or(DateTime::::UNIX_EPOCH), + headers: headers.unwrap_or_default(), + user_metadata: user_metadata.unwrap_or_default(), + system_metadata: system_metadata.unwrap_or_default(), } } fn __repr__(&self) -> String { format!( - "Datum(keys={:?}, value={:?}, watermark={}, eventtime={}, headers={:?}, user_metadata={:?}, system_metadata={:?})", + "Datum(keys={:?}, value={}, watermark={}, event_time={}, headers={:?}, user_metadata={}, system_metadata={})", self.keys, - self.value, + bytes_literal(&self.value), self.watermark, - self.eventtime, + self.event_time, self.headers, - self.user_metadata, - self.system_metadata + metadata_literal(&self.user_metadata), + metadata_literal(&self.system_metadata) ) } fn __str__(&self) -> String { - format!( - "Datum(keys={:?}, value={:?}, watermark={}, eventtime={}, headers={:?}, user_metadata={:?}, system_metadata={:?})", - self.keys, - String::from_utf8_lossy(&self.value), - self.watermark, - self.eventtime, - self.headers, - self.user_metadata, - self.system_metadata - ) + self.__repr__() } } impl From for Datum { fn from(value: map::MapRequest) -> Self { - Datum::new( - value.keys, - value.value, - value.watermark, - value.eventtime, - value.headers, - value.user_metadata.into(), - value.system_metadata.into(), - ) + Self { + keys: value.keys, + value: value.value, + watermark: value.watermark, + event_time: value.eventtime, + headers: value.headers, + user_metadata: user_metadata_to_hash_map(value.user_metadata), + system_metadata: system_metadata_to_hash_map(value.system_metadata), + } + } +} + +pub(crate) async fn map_grpc_client( + sock_file: String, +) -> PyResult> { + let endpoint = tonic::transport::Endpoint::try_from("http://[::]:50051") + .map_err(|e| pyo3::PyErr::new::(e.to_string()))?; + + let channel = endpoint + .connect_with_connector(service_fn(move |_: Uri| { + let sock = PathBuf::from(sock_file.clone()); + async move { + Ok::<_, std::io::Error>(hyper_util::rt::TokioIo::new( + UnixStream::connect(sock).await?, + )) + } + })) + .await + .map_err(|e| pyo3::PyErr::new::(e.to_string()))?; + + Ok(MapClient::new(channel)) +} + +pub(crate) async fn wait_for_ready( + sock_file: String, + timeout: Duration, + component: &'static str, +) -> PyResult<()> { + let deadline = Instant::now() + timeout; + + loop { + if let Ok(mut client) = map_grpc_client(sock_file.clone()).await + && let Ok(response) = client.is_ready(()).await + && response.into_inner().ready + { + return Ok(()); + } + + let now = Instant::now(); + if now >= deadline { + return Err(pyo3::PyErr::new::( + format!("timed out waiting for {component} server readiness"), + )); + } + + tokio::time::sleep(std::cmp::min( + Duration::from_millis(100), + deadline.saturating_duration_since(now), + )) + .await; } } /// Async Map Server that can be started from Python code which will run the Python UDF function. -#[pyclass(module = "pynumaflow_lite.mapper")] +#[pyclass(name = "_MapAsyncServer", module = "pynumaflow_lite.mapper")] pub struct MapAsyncServer { sock_file: String, - info_file: String, + server_info_file: String, shutdown_tx: Mutex>>, } #[pymethods] impl MapAsyncServer { #[new] - #[pyo3(signature = (sock_file: "str | None"=map::SOCK_ADDR.to_string(), info_file: "str | None"=map::SERVER_INFO_FILE.to_string()) -> "MapAsyncServer" - )] - fn new(sock_file: String, info_file: String) -> Self { + #[pyo3(signature = ( + sock_file: "str | None"=None, + server_info_file: "str | None"=None, + ) -> "_MapAsyncServer")] + fn new(sock_file: Option, server_info_file: Option) -> Self { Self { - sock_file, - info_file, + sock_file: sock_file.unwrap_or_else(|| map::SOCK_ADDR.to_string()), + server_info_file: server_info_file.unwrap_or_else(|| map::SERVER_INFO_FILE.to_string()), 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> { + #[pyo3(signature = (handler: "callable") -> "None")] + pub fn start<'a>(&self, py: Python<'a>, handler: Py) -> PyResult> { let sock_file = self.sock_file.clone(); - let info_file = self.info_file.clone(); + let server_info_file = self.server_info_file.clone(); let (tx, rx) = tokio::sync::oneshot::channel::<()>(); { let mut guard = self.shutdown_tx.lock().unwrap(); @@ -390,9 +330,25 @@ impl MapAsyncServer { } pyo3_async_runtimes::tokio::future_into_py(py, async move { - crate::map::server::start(py_func, sock_file, info_file, rx) - .await - .expect("server failed to start"); + crate::map::server::start(handler, sock_file, server_info_file, rx).await?; + Ok(()) + }) + } + + /// Wait until the Numaflow IsReady probe succeeds over the map UDS. + #[pyo3(signature = (timeout: "float"=30.0) -> "None")] + pub fn wait_ready<'a>(&self, py: Python<'a>, timeout: f64) -> PyResult> { + if !timeout.is_finite() || timeout < 0.0 { + return Err(pyo3::PyErr::new::( + "timeout must be a non-negative finite float", + )); + } + + let sock_file = self.sock_file.clone(); + let timeout = Duration::from_secs_f64(timeout); + + pyo3_async_runtimes::tokio::future_into_py(py, async move { + wait_for_ready(sock_file, timeout, "map").await?; Ok(()) }) } @@ -409,9 +365,6 @@ impl MapAsyncServer { /// Helper to populate a PyModule with map 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::()?; diff --git a/packages/pynumaflow-lite/src/map/server.rs b/packages/pynumaflow-lite/src/map/server.rs index 1e220994..60d7d83e 100644 --- a/packages/pynumaflow-lite/src/map/server.rs +++ b/packages/pynumaflow-lite/src/map/server.rs @@ -1,40 +1,114 @@ -use crate::map::{Datum, Messages}; use numaflow::map; use numaflow::shared::ServerExtras; +use pyo3::exceptions::PyTypeError; use pyo3::prelude::*; -use std::sync::Arc; +use std::sync::{Arc, Mutex}; pub(crate) struct PyMapRunner { pub(crate) event_loop: Arc>, pub(crate) py_func: Arc>, + pub(crate) errors: Arc>>, } -#[tonic::async_trait] -impl map::Mapper for PyMapRunner { - async fn map(&self, input: map::MapRequest) -> Vec { - let fut = Python::attach(|py| { - let keys = input.keys.clone(); - let input: Datum = input.into(); - let py_func = self.py_func.clone(); +// Build the full Python traceback text for the panic message, so the sidecar +// reports the same failure that Python raises. +fn format_error(py: Python<'_>, error: &PyErr) -> String { + match error.traceback(py).map(|traceback| traceback.format()) { + Some(Ok(traceback)) => format!("{traceback}{error}"), + _ => error.to_string(), + } +} - let locals = pyo3_async_runtimes::TaskLocals::new(self.event_loop.bind(py).clone()); +// Join every handler failure into one error for Python to raise. +// +// Python 3.11 and later have BaseExceptionGroup, which prints each traceback in +// turn. Older versions have no group type, so they get the first error only. +fn combine_errors(py: Python<'_>, errors: Vec) -> PyErr { + let first = || errors.first().expect("errors is never empty").clone_ref(py); - let coro = py_func.call1(py, (keys, input)).unwrap().into_bound(py); + if errors.len() == 1 { + return first(); + } - pyo3_async_runtimes::into_future_with_locals(&locals, coro).unwrap() - }); + let Ok(group_type) = py + .import("builtins") + .and_then(|builtins| builtins.getattr("BaseExceptionGroup")) + else { + return first(); + }; - let result = fut.await.unwrap(); + let values: Vec<_> = errors.iter().map(|error| error.value(py).clone()).collect(); + let message = format!("{} map handler calls failed", values.len()); - let result = Python::attach(|py| { - let x: Messages = result.extract(py).unwrap(); - x - }); + match group_type.call1((message, values)) { + Ok(group) => PyErr::from_value(group), + Err(_) => first(), + } +} - println!("{:?}", result); +impl PyMapRunner { + fn fail(&self, error: PyErr) -> ! { + // numaflow calls map() concurrently, so each error belongs to a different + // message. Keep all of them. start() raises them together, which lets + // Python format every traceback instead of Rust printing them by hand. + let message = Python::attach(|py| format_error(py, &error)); + self.errors.lock().unwrap().push(error); + + // numaflow catches this panic, sends a gRPC error for this message, and + // starts the server shutdown. An empty result would instead look like a + // message that the handler dropped on purpose. + panic!("{message}"); + } +} - result.messages.into_iter().map(|m| m.into()).collect() +#[tonic::async_trait] +impl map::Mapper for PyMapRunner { + async fn map(&self, input: map::MapRequest) -> Vec { + let fut = match Python::attach(|py| -> PyResult<_> { + let locals = pyo3_async_runtimes::TaskLocals::new(self.event_loop.bind(py).clone()); + let datum: crate::map::Datum = input.into(); + let coro = self.py_func.call1(py, (datum,))?.into_bound(py); + let is_awaitable: bool = py + .import("inspect")? + .call_method1("isawaitable", (&coro,))? + .extract()?; + if !is_awaitable { + return Err(PyErr::new::( + "map handler must be an async function (coroutine)", + )); + } + pyo3_async_runtimes::into_future_with_locals(&locals, coro).map_err(|_| { + PyErr::new::("map handler must be an async function (coroutine)") + }) + }) { + Ok(fut) => fut, + Err(error) => self.fail(error), + }; + + let result = match fut.await { + Ok(result) => result, + Err(error) => self.fail(error), + }; + + let messages: Vec = match Python::attach(|py| { + result.extract(py).map_err(|_| { + let type_name = result + .bind(py) + .get_type() + .name() + .map(|name| name.to_string_lossy().into_owned()) + .unwrap_or_else(|_| "".to_string()); + PyErr::new::(format!( + "map handler must return list[Message], got {type_name}" + )) + }) + }) { + Ok(messages) => messages, + Err(error) => self.fail(error), + }; + + messages.into_iter().map(|m| m.into()).collect() } } @@ -46,14 +120,24 @@ pub(super) async fn start( 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 py_asyncio_loop_handle = tokio::task::spawn_blocking({ + println!( + "Starting Map UDF. socket={}, server_info={}", + sock_file, info_file + ); + 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 errors = Arc::new(Mutex::new(Vec::new())); + // Shutdown has two sources, and neither one needs a channel here. The Python + // side signals stop() through shutdown_rx. An uncaught Python error panics in + // fail(), and numaflow then shuts the server down on its own. let py_map_runner = PyMapRunner { py_func: Arc::new(py_func), event_loop: event_loop.clone(), + errors: errors.clone(), }; let server = numaflow::map::Server::new(py_map_runner) @@ -61,7 +145,7 @@ pub(super) async fn start( .with_server_info_file(info_file); let result = server - .start_with_shutdown(combined_rx) + .start_with_shutdown(shutdown_rx) .await .map_err(|e| pyo3::PyErr::new::(e.to_string())); @@ -72,15 +156,14 @@ pub(super) async fn start( } }); - println!("Numaflow Core has shutdown..."); + println!("Numaflow Map 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"); - sig_handle.abort(); + let errors = std::mem::take(&mut *errors.lock().unwrap()); + if !errors.is_empty() { + return Err(Python::attach(|py| combine_errors(py, errors))); } result diff --git a/packages/pynumaflow-lite/tests/examples/map_cat.py b/packages/pynumaflow-lite/tests/examples/map_cat.py index 9e04728b..bb38d770 100644 --- a/packages/pynumaflow-lite/tests/examples/map_cat.py +++ b/packages/pynumaflow-lite/tests/examples/map_cat.py @@ -1,61 +1,34 @@ -import asyncio -import signal -from collections.abc import Awaitable, Callable +from pynumaflow_lite.mapper import Datum, MapAsyncServer, Message -from pynumaflow_lite import mapper - - -async def async_handler(keys: list[str], payload: mapper.Datum) -> mapper.Messages: - messages = mapper.Messages() +async def map_handler(datum: Datum) -> list[Message]: # Read system metadata (read-only) - print(f"System metadata groups: {payload.system_metadata.groups()}") - for group in payload.system_metadata.groups(): - for key in payload.system_metadata.keys(group): - value = payload.system_metadata.value(group, key) + print(f"System metadata groups: {list(datum.system_metadata)}") + for group, key_values in datum.system_metadata.items(): + for key, value in key_values.items(): print(f" System[{group}][{key}] = {value}") - # Read user metadata (read-only from input) - print(f"User metadata groups: {payload.user_metadata.groups()}") - for group in payload.user_metadata.groups(): - for key in payload.user_metadata.keys(group): - value = payload.user_metadata.value(group, key) + # Read user metadata + print(f"User metadata groups: {list(datum.user_metadata)}") + for group, key_values in datum.user_metadata.items(): + for key, value in key_values.items(): print(f" User[{group}][{key}] = {value}") - if payload.value == b"bad world": - messages.append(mapper.Message.message_to_drop()) - else: - # Create user metadata for the outgoing message - user_metadata = mapper.UserMetadata() - user_metadata.create_group("processing") - user_metadata.add_kv("processing", "handler", b"map_cat") - user_metadata.add_kv("processing", "msg_length", str(len(payload.value)).encode()) - - messages.append(mapper.Message(payload.value, keys, user_metadata=user_metadata)) - - return messages - - -async def start(f: Callable[[list[str], mapper.Datum], Awaitable[mapper.Messages]]): - sock_file = "/tmp/var/run/numaflow/map.sock" - server_info_file = "/tmp/var/run/numaflow/mapper-server-info" - server = mapper.MapAsyncServer(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 + if datum.value == b"bad world": + return [Message.to_drop()] - try: - await server.start(f) - print("Shutting down gracefully...") - except asyncio.CancelledError: - server.stop() - return + user_metadata = { + "processing": { + "handler": b"map_cat", + "msg_length": str(len(datum.value)).encode(), + } + } + return [Message(datum.value, keys=datum.keys, user_metadata=user_metadata)] if __name__ == "__main__": - asyncio.run(start(async_handler)) + MapAsyncServer( + map_handler, + sock_file="/tmp/var/run/numaflow/map.sock", + server_info_file="/tmp/var/run/numaflow/mapper-server-info", + ).run() diff --git a/packages/pynumaflow-lite/tests/examples/map_cat_class.py b/packages/pynumaflow-lite/tests/examples/map_cat_class.py index 079beaf2..fba24163 100644 --- a/packages/pynumaflow-lite/tests/examples/map_cat_class.py +++ b/packages/pynumaflow-lite/tests/examples/map_cat_class.py @@ -1,70 +1,35 @@ -import asyncio -import signal -from collections.abc import Awaitable, Callable +from pynumaflow_lite.mapper import Datum, MapAsyncServer, Mapper, Message -from pynumaflow_lite import mapper - - -class SimpleCat(mapper.Mapper): - async def handler(self, keys: list[str], payload: mapper.Datum) -> mapper.Messages: - - messages = mapper.Messages() +class SimpleCat(Mapper): + async def handler(self, datum: Datum) -> list[Message]: # Read system metadata (read-only) - print(f"System metadata groups: {payload.system_metadata.groups()}") - for group in payload.system_metadata.groups(): - for key in payload.system_metadata.keys(group): - value = payload.system_metadata.value(group, key) + print(f"System metadata groups: {list(datum.system_metadata)}") + for group, key_values in datum.system_metadata.items(): + for key, value in key_values.items(): print(f" System[{group}][{key}] = {value}") - # Read user metadata (read-only from input) - print(f"User metadata groups: {payload.user_metadata.groups()}") - for group in payload.user_metadata.groups(): - for key in payload.user_metadata.keys(group): - value = payload.user_metadata.value(group, key) + # Read user metadata + print(f"User metadata groups: {list(datum.user_metadata)}") + for group, key_values in datum.user_metadata.items(): + for key, value in key_values.items(): print(f" User[{group}][{key}] = {value}") - if payload.value == b"bad world": - messages.append(mapper.Message.message_to_drop()) - else: - # Create user metadata for the outgoing message - user_metadata = mapper.UserMetadata() - user_metadata.create_group("processing") - user_metadata.add_kv("processing", "handler", b"map_cat_class") - user_metadata.add_kv("processing", "msg_length", str(len(payload.value)).encode()) - - messages.append(mapper.Message(payload.value, keys, user_metadata=user_metadata)) - - return messages - - -async def start(f: Callable[[list[str], mapper.Datum], Awaitable[mapper.Messages]]): - sock_file = "/tmp/var/run/numaflow/map.sock" - server_info_file = "/tmp/var/run/numaflow/mapper-server-info" - server = mapper.MapAsyncServer(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 + if datum.value == b"bad world": + return [Message.to_drop()] - try: - await server.start(f) - print("Shutting down gracefully...") - except asyncio.CancelledError: - # Fallback in case the task was cancelled by the runner - server.stop() - return + user_metadata = { + "processing": { + "handler": b"map_cat_class", + "msg_length": str(len(datum.value)).encode(), + } + } + return [Message(datum.value, keys=datum.keys, user_metadata=user_metadata)] if __name__ == "__main__": - async_handler = SimpleCat() - asyncio.run(start(async_handler)) + MapAsyncServer( + SimpleCat(), + sock_file="/tmp/var/run/numaflow/map.sock", + server_info_file="/tmp/var/run/numaflow/mapper-server-info", + ).run()