diff --git a/packages/pynumaflow-lite/manifests/batchmap/.dockerignore b/packages/pynumaflow-lite/manifests/batchmap/.dockerignore
new file mode 100644
index 00000000..21d0b898
--- /dev/null
+++ b/packages/pynumaflow-lite/manifests/batchmap/.dockerignore
@@ -0,0 +1 @@
+.venv/
diff --git a/packages/pynumaflow-lite/manifests/batchmap/Dockerfile b/packages/pynumaflow-lite/manifests/batchmap/Dockerfile
index af12e945..5caaa1a1 100644
--- a/packages/pynumaflow-lite/manifests/batchmap/Dockerfile
+++ b/packages/pynumaflow-lite/manifests/batchmap/Dockerfile
@@ -1,37 +1,43 @@
-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", "batchmap_cat.py"]
\ No newline at end of file
+FROM python:3.11-slim-trixie AS builder
+
+COPY --from=ghcr.io/astral-sh/uv:0.12.13 /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
+
+# Install the local pynumaflow-lite wheel that matches the build platform.
+RUN uv pip install --no-index --find-links . pynumaflow-lite
+
+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", "batchmap_cat.py"]
diff --git a/packages/pynumaflow-lite/manifests/batchmap/batchmap_cat.py b/packages/pynumaflow-lite/manifests/batchmap/batchmap_cat.py
index fed9eedf..e6fd73b6 100644
--- a/packages/pynumaflow-lite/manifests/batchmap/batchmap_cat.py
+++ b/packages/pynumaflow-lite/manifests/batchmap/batchmap_cat.py
@@ -1,46 +1,26 @@
 import asyncio
-import signal
-from collections.abc import AsyncIterable, Awaitable, Callable
+from collections.abc import AsyncIterable
 
-from pynumaflow_lite import batchmapper
-from pynumaflow_lite.batchmapper import Message
+from pynumaflow_lite.batchmapper import BatchMapAsyncServer, BatchMapper, BatchResponse, Datum, Message
 
 
-class SimpleBatchCat(batchmapper.BatchMapper):
-    async def handler(self, batch: AsyncIterable[batchmapper.Datum]) -> batchmapper.BatchResponses:
-        responses = batchmapper.BatchResponses()
-        async for d in batch:
-            resp = batchmapper.BatchResponse(d.id)
-            if d.value == b"bad world":
-                resp.append(Message.message_to_drop())
-                continue
-
-            resp.append(Message(d.value, d.keys))
-            responses.append(resp)
+class SimpleBatchCat(BatchMapper):
+    async def handler(self, batch: AsyncIterable[Datum]) -> list[BatchResponse]:
+        responses = []
+        async for datum in batch:
+            if datum.value == b"bad world":
+                responses.append(BatchResponse(datum.id, Message.to_drop()))
+            else:
+                responses.append(BatchResponse(datum.id, Message(datum.value, keys=datum.keys)))
         return responses
 
 
-async def start(
-    f: Callable[[AsyncIterable[batchmapper.Datum]], Awaitable[batchmapper.BatchResponses]],
-):
-    server = batchmapper.BatchMapAsyncServer()
-
-    # Register loop-level signal handlers so we control shutdown and avoid asyncio.run
-    loop = asyncio.get_running_loop()
-    try:
-        loop.add_signal_handler(signal.SIGINT, lambda: server.stop())
-        loop.add_signal_handler(signal.SIGTERM, lambda: server.stop())
-    except (NotImplementedError, RuntimeError):
-        pass
-
-    try:
-        await server.start(f)
-        print("Shutting down gracefully...")
-    except asyncio.CancelledError:
-        server.stop()
-        return
+async def main() -> None:
+    print("Starting BatchMap server")
+    # `serve` returns when SIGINT or SIGTERM arrives.
+    await BatchMapAsyncServer(SimpleBatchCat()).serve()
+    print("BatchMap server stopped")
 
 
 if __name__ == "__main__":
-    async_handler = SimpleBatchCat()
-    asyncio.run(start(async_handler))
+    asyncio.run(main())
diff --git a/packages/pynumaflow-lite/manifests/batchmap/pipeline.yaml b/packages/pynumaflow-lite/manifests/batchmap/pipeline.yaml
index 3f3372b7..f179a829 100644
--- a/packages/pynumaflow-lite/manifests/batchmap/pipeline.yaml
+++ b/packages/pynumaflow-lite/manifests/batchmap/pipeline.yaml
@@ -8,8 +8,8 @@ spec:
       source:
         # A self data generating source
         generator:
-          rpu: 500
-          duration: 1s
+          rpu: 3
+          duration: 3s
     - name: batchmap
       partitions: 2
       scale:
@@ -27,4 +27,4 @@ spec:
     - from: in
       to: batchmap
     - from: batchmap
-      to: sink
\ No newline at end of file
+      to: sink
diff --git a/packages/pynumaflow-lite/manifests/batchmap/pyproject.toml b/packages/pynumaflow-lite/manifests/batchmap/pyproject.toml
index de7ee701..1d74d87b 100644
--- a/packages/pynumaflow-lite/manifests/batchmap/pyproject.toml
+++ b/packages/pynumaflow-lite/manifests/batchmap/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/batchmap/uv.lock b/packages/pynumaflow-lite/manifests/batchmap/uv.lock
new file mode 100644
index 00000000..c9ded2b3
--- /dev/null
+++ b/packages/pynumaflow-lite/manifests/batchmap/uv.lock
@@ -0,0 +1,8 @@
+version = 1
+revision = 3
+requires-python = "==3.11.*"
+
+[[package]]
+name = "batchmap-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 46a781e2..c346b880 100644
--- a/packages/pynumaflow-lite/pynumaflow_lite/__init__.py
+++ b/packages/pynumaflow-lite/pynumaflow_lite/__init__.py
@@ -1,69 +1,11 @@
 from . import (
     pynumaflow_lite,  # type: ignore[attr-defined]  # Rust extension, resolved at runtime
 )
-from .pynumaflow_lite import *  # noqa: F403  # Rust extension; exports resolved at runtime
-
-# Ensure the `mapper`, `batchmapper`, and `mapstreamer` submodules are importable as attributes of the package
-# even though they're primarily registered by the extension module.
-try:
-    from importlib import import_module as _import_module
-
-    mapper = _import_module(__name__ + ".mapper")
-except Exception:  # pragma: no cover - avoid hard failures if extension not built
-    mapper = None
-
-try:
-    batchmapper = _import_module(__name__ + ".batchmapper")
-except Exception:  # pragma: no cover
-    batchmapper = None
-
-try:
-    mapstreamer = _import_module(__name__ + ".mapstreamer")
-except Exception:  # pragma: no cover
-    mapstreamer = None
-try:
-    reducer = _import_module(__name__ + ".reducer")
-except Exception:  # pragma: no cover
-    reducer = None
-
-try:
-    session_reducer = _import_module(__name__ + ".session_reducer")
-except Exception:  # pragma: no cover
-    session_reducer = None
-
-try:
-    reducestreamer = _import_module(__name__ + ".reducestreamer")
-except Exception:  # pragma: no cover
-    reducestreamer = None
-
-try:
-    accumulator = _import_module(__name__ + ".accumulator")
-except Exception:  # pragma: no cover
-    accumulator = None
-
-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
-
-try:
-    sourcetransformer = _import_module(__name__ + ".sourcetransformer")
-except Exception:  # pragma: no cover
-    sourcetransformer = None
-
-try:
-    sideinputer = _import_module(__name__ + ".sideinputer")
-except Exception:  # pragma: no cover
-    sideinputer = None
 
 # Surface the Python Mapper, BatchMapper, MapStreamer, Reducer, SessionReducer, ReduceStreamer, Accumulator, Sinker,
 # Sourcer, SourceTransformer, and SideInput classes under the extension submodules for convenient access
 from ._accumulator_dtypes import Accumulator
+from ._batchmap_server import BatchMapAsyncServer
 from ._batchmapper_dtypes import BatchMapper
 from ._map_dtypes import Mapper
 from ._map_server import MapAsyncServer
@@ -76,41 +18,38 @@
 from ._sink_server import SinkAsyncServer
 from ._source_dtypes import Sourcer
 from ._sourcetransformer_dtypes import SourceTransformer
+from .pynumaflow_lite import *  # noqa: F403  # Rust extension; exports resolved at runtime
 
-if mapper is not None:
-    mapper.Mapper = Mapper
-    mapper.MapAsyncServer = MapAsyncServer
-
-if batchmapper is not None:
-    batchmapper.BatchMapper = BatchMapper
-
-if mapstreamer is not None:
-    mapstreamer.MapStreamer = MapStreamer
-
-if reducer is not None:
-    reducer.Reducer = Reducer
-
-if session_reducer is not None:
-    session_reducer.SessionReducer = SessionReducer
-
-if reducestreamer is not None:
-    reducestreamer.ReduceStreamer = ReduceStreamer
-
-if accumulator is not None:
-    accumulator.Accumulator = Accumulator
-
-if sinker is not None:
-    sinker.Sinker = Sinker
-    sinker.SinkAsyncServer = SinkAsyncServer
-
-if sourcer is not None:
-    sourcer.Sourcer = Sourcer
-
-if sourcetransformer is not None:
-    sourcetransformer.SourceTransformer = SourceTransformer
+# Submodules are defined by the Rust extension, which also registers them in sys.modules
+# as `pynumaflow_lite.<name>`.
+from .pynumaflow_lite import (  # type: ignore[attr-defined]
+    accumulator,
+    batchmapper,
+    mapper,
+    mapstreamer,
+    reducer,
+    reducestreamer,
+    session_reducer,
+    sideinputer,
+    sinker,
+    sourcer,
+    sourcetransformer,
+)
 
-if sideinputer is not None:
-    sideinputer.SideInput = SideInput
+mapper.Mapper = Mapper
+mapper.MapAsyncServer = MapAsyncServer
+batchmapper.BatchMapper = BatchMapper
+batchmapper.BatchMapAsyncServer = BatchMapAsyncServer
+mapstreamer.MapStreamer = MapStreamer
+reducer.Reducer = Reducer
+session_reducer.SessionReducer = SessionReducer
+reducestreamer.ReduceStreamer = ReduceStreamer
+accumulator.Accumulator = Accumulator
+sinker.Sinker = Sinker
+sinker.SinkAsyncServer = SinkAsyncServer
+sourcer.Sourcer = Sourcer
+sourcetransformer.SourceTransformer = SourceTransformer
+sideinputer.SideInput = SideInput
 
 # Public API
 __all__ = [
@@ -128,6 +67,3 @@
 ]
 
 __doc__ = pynumaflow_lite.__doc__
-if hasattr(pynumaflow_lite, "__all__"):
-    # Merge to keep our package-level exports
-    __all__ = list(set(__all__) | set(pynumaflow_lite.__all__))
diff --git a/packages/pynumaflow-lite/pynumaflow_lite/_batchmap_server.py b/packages/pynumaflow-lite/pynumaflow_lite/_batchmap_server.py
new file mode 100644
index 00000000..32b15304
--- /dev/null
+++ b/packages/pynumaflow-lite/pynumaflow_lite/_batchmap_server.py
@@ -0,0 +1,132 @@
+from __future__ import annotations
+
+import asyncio
+import contextlib
+import signal
+from collections.abc import AsyncIterable, Awaitable, Callable
+from types import TracebackType
+from typing import TypeAlias
+
+from .pynumaflow_lite import batchmapper as _batchmapper
+
+Datum: TypeAlias = _batchmapper.Datum
+BatchResponse: TypeAlias = _batchmapper.BatchResponse
+
+_SHUTDOWN_SIGNALS = (signal.SIGINT, signal.SIGTERM)
+
+
+class BatchMapAsyncServer:
+    def __init__(
+        self,
+        handler: Callable[[AsyncIterable[Datum]], Awaitable[list[BatchResponse]]],
+        *,
+        sock_file: str | None = None,
+        server_info_file: str | None = None,
+        install_signal_handlers: bool = True,
+    ):
+        self._core = _batchmapper._BatchMapAsyncServer(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 batchmap 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("batchmap 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("batchmap 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) -> BatchMapAsyncServer:
+        """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("batchmap 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 batchmap 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/_map_dtypes.py b/packages/pynumaflow-lite/pynumaflow_lite/_map_dtypes.py
index cb92c551..5ceb0026 100644
--- a/packages/pynumaflow-lite/pynumaflow_lite/_map_dtypes.py
+++ b/packages/pynumaflow-lite/pynumaflow_lite/_map_dtypes.py
@@ -17,7 +17,7 @@ class instance is sent as a callable.
         return self.handler(*args, **kwargs)
 
     @abstractmethod
-    async def handler(self, payload: Datum) -> list[Message]:
+    async def handler(self, datum: 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
index c758051b..ff477680 100644
--- a/packages/pynumaflow-lite/pynumaflow_lite/_map_server.py
+++ b/packages/pynumaflow-lite/pynumaflow_lite/_map_server.py
@@ -5,11 +5,12 @@
 import signal
 from collections.abc import Awaitable, Callable
 from types import TracebackType
+from typing import TypeAlias
 
 from .pynumaflow_lite import mapper as _mapper
 
-Datum = _mapper.Datum
-Message = _mapper.Message
+Datum: TypeAlias = _mapper.Datum
+Message: TypeAlias = _mapper.Message
 
 _SHUTDOWN_SIGNALS = (signal.SIGINT, signal.SIGTERM)
 
diff --git a/packages/pynumaflow-lite/pynumaflow_lite/batchmapper.pyi b/packages/pynumaflow-lite/pynumaflow_lite/batchmapper.pyi
index 1fb96f12..a60d81b5 100644
--- a/packages/pynumaflow-lite/pynumaflow_lite/batchmapper.pyi
+++ b/packages/pynumaflow-lite/pynumaflow_lite/batchmapper.pyi
@@ -1,7 +1,10 @@
 from __future__ import annotations
 
 import datetime as _dt
-from collections.abc import AsyncIterator, Awaitable, Callable
+from collections.abc import AsyncIterable, Awaitable, Callable
+from types import TracebackType
+
+from ._batchmapper_dtypes import BatchMapper as BatchMapper
 
 class NackOptions:
     """Per-message redelivery options for a nack."""
@@ -32,53 +35,80 @@ class Message:
         keys: list[str] | None = ...,
         tags: list[str] | None = ...,
     ) -> None: ...
+    def __repr__(self) -> str: ...
+    def __eq__(self, other: object) -> bool: ...
     @staticmethod
-    def message_to_drop() -> Message: ...
+    def to_drop() -> Message: ...
     @staticmethod
     def to_nack(nack_options: NackOptions | None = ...) -> Message: ...
     @staticmethod
     def to_fail() -> Message: ...
 
 class Datum:
+    id: str
     keys: list[str]
     value: bytes
     watermark: _dt.datetime
-    eventtime: _dt.datetime
-    id: str
+    event_time: _dt.datetime
     headers: dict[str, str]
 
+    def __init__(
+        self,
+        *,
+        id: str,
+        keys: list[str] | None = ...,
+        value: bytes | None = ...,
+        event_time: _dt.datetime | None = ...,
+        watermark: _dt.datetime | None = ...,
+        headers: dict[str, str] | None = ...,
+    ) -> None: ...
     def __repr__(self) -> str: ...
     def __str__(self) -> str: ...
 
 class BatchResponse:
     id: str
+    messages: list[Message]
 
-    def __init__(self, id: str) -> None: ...
-    @staticmethod
-    def from_id(id: str) -> BatchResponse: ...
+    def __init__(self, id: str, *messages: Message) -> None: ...
     def append(self, message: Message) -> None: ...
+    def __repr__(self) -> str: ...
 
-class BatchResponses:
-    def __init__(self) -> None: ...
-    def append(self, response: BatchResponse) -> None: ...
+class _BatchMapAsyncServer:
+    def __init__(
+        self,
+        sock_file: str | None = ...,
+        server_info_file: str | None = ...,
+    ) -> None: ...
+    def start(self, handler: Callable[[AsyncIterable[Datum]], Awaitable[list[BatchResponse]]]) -> Awaitable[None]: ...
+    def wait_ready(self, timeout: float = ...) -> Awaitable[None]: ...
+    def stop(self) -> None: ...
 
 class BatchMapAsyncServer:
     def __init__(
         self,
+        handler: Callable[[AsyncIterable[Datum]], Awaitable[list[BatchResponse]]],
+        *,
         sock_file: str | None = ...,
-        info_file: str | None = ...,
+        server_info_file: str | None = ...,
+        install_signal_handlers: bool = ...,
     ) -> None: ...
-    def start(self, py_func: Callable[[AsyncIterator[Datum]], Awaitable[BatchResponses]]) -> Awaitable[None]: ...
+    def run(self) -> None: ...
+    async def serve(self) -> None: ...
     def stop(self) -> None: ...
-
-class BatchMapper:
-    async def handler(self, batch: AsyncIterator[Datum]) -> BatchResponses: ...
+    async def wait_ready(self, timeout: float = ...) -> None: ...
+    async def wait_for_termination(self) -> None: ...
+    async def __aenter__(self) -> BatchMapAsyncServer: ...
+    async def __aexit__(
+        self,
+        exc_type: type[BaseException] | None,
+        exc: BaseException | None,
+        tb: TracebackType | None,
+    ) -> None: ...
 
 __all__ = [
     "BatchMapAsyncServer",
     "BatchMapper",
     "BatchResponse",
-    "BatchResponses",
     "Datum",
     "Message",
     "NackOptions",
diff --git a/packages/pynumaflow-lite/src/batchmap/mod.rs b/packages/pynumaflow-lite/src/batchmap/mod.rs
index ef929b78..3c73b253 100644
--- a/packages/pynumaflow-lite/src/batchmap/mod.rs
+++ b/packages/pynumaflow-lite/src/batchmap/mod.rs
@@ -1,4 +1,4 @@
-use std::collections::HashMap;
+use std::{collections::HashMap, time::Duration};
 
 use numaflow::batchmap;
 
@@ -14,17 +14,21 @@ use pyo3::prelude::*;
 use std::sync::Mutex;
 
 use crate::nack::NackOptions;
+use crate::pyrs::{bytes_literal, py_repr};
 
 /// A message to be sent for a single datum in batch response.
-#[pyclass(module = "pynumaflow_lite.batchmapper", from_py_object)]
-#[derive(Clone, Default, Debug)]
+#[pyclass(module = "pynumaflow_lite.batchmapper", 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<Vec<String>>,
     /// Value is the value passed to the next vertex.
+    #[pyo3(get)]
     pub value: Vec<u8>,
     /// Tags are used for conditional forwarding.
+    #[pyo3(get)]
     pub tags: Option<Vec<String>>,
     /// Options sent back to the source when nacking the message.
     #[pyo3(get)]
@@ -33,7 +37,7 @@ pub struct Message {
 
 #[pymethods]
 impl Message {
-    /// Create a new [Message] with the given value, keys, and tags.
+    /// Create a new [Message] with the given value. The keys and tags are optional.
     #[new]
     #[pyo3(signature = (value: "bytes", keys: "list[str] | None"=None, tags: "list[str] | None"=None) -> "Message"
     )]
@@ -47,9 +51,9 @@ impl Message {
     }
 
     /// Drop a [Message], do not forward to the next vertex.
-    #[pyo3(signature = ())]
+    #[pyo3(signature = () -> "Message")]
     #[staticmethod]
-    fn message_to_drop() -> Self {
+    fn to_drop() -> Self {
         Self {
             keys: None,
             value: vec![],
@@ -81,6 +85,18 @@ impl Message {
             nack_options: None,
         }
     }
+
+    fn __repr__(&self, py: Python<'_>) -> PyResult<String> {
+        Ok(format!(
+            "Message(value={}, keys={}, tags={}, nack_options={})",
+            bytes_literal(&self.value),
+            py_repr(py, &self.keys)?,
+            py_repr(py, &self.tags)?,
+            self.nack_options
+                .as_ref()
+                .map_or_else(|| "None".to_string(), NackOptions::__repr__),
+        ))
+    }
 }
 
 impl From<Message> for batchmap::Message {
@@ -94,7 +110,7 @@ impl From<Message> for batchmap::Message {
     }
 }
 
-/// The incoming Datum for BatchMap
+/// The incoming Datum passed to the BatchMap handler.
 #[pyclass(module = "pynumaflow_lite.batchmapper")]
 pub struct Datum {
     /// Set of keys in the (key, value) terminology of map/reduce paradigm.
@@ -108,7 +124,7 @@ pub struct Datum {
     pub watermark: DateTime<Utc>,
     /// Time of the element as seen at source or aligned after a reduce operation.
     #[pyo3(get)]
-    pub eventtime: DateTime<Utc>,
+    pub event_time: DateTime<Utc>,
     /// ID is the unique id of the message
     #[pyo3(get)]
     pub id: String,
@@ -117,54 +133,62 @@ pub struct Datum {
     pub headers: HashMap<String, String>,
 }
 
+#[pymethods]
 impl Datum {
+    #[new]
+    #[pyo3(signature = (
+        *,
+        id: "str",
+        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,
+    ) -> "Datum")]
     fn new(
-        keys: Vec<String>,
-        value: Vec<u8>,
-        watermark: DateTime<Utc>,
-        eventtime: DateTime<Utc>,
         id: String,
-        headers: HashMap<String, String>,
+        keys: Option<Vec<String>>,
+        value: Option<Vec<u8>>,
+        event_time: Option<DateTime<Utc>>,
+        watermark: Option<DateTime<Utc>>,
+        headers: Option<HashMap<String, String>>,
     ) -> Self {
         Self {
-            keys,
-            value,
-            watermark,
-            eventtime,
             id,
-            headers,
+            keys: keys.unwrap_or_default(),
+            value: value.unwrap_or_default(),
+            watermark: watermark.unwrap_or(DateTime::<Utc>::UNIX_EPOCH),
+            event_time: event_time.unwrap_or(DateTime::<Utc>::UNIX_EPOCH),
+            headers: headers.unwrap_or_default(),
         }
     }
 
-    fn __repr__(&self) -> String {
-        format!(
-            "Datum(keys={:?}, value={:?}, watermark={}, eventtime={}, id={}, headers={:?})",
-            self.keys, self.value, self.watermark, self.eventtime, self.id, self.headers
-        )
+    fn __repr__(&self, py: Python<'_>) -> PyResult<String> {
+        Ok(format!(
+            "Datum(id={}, keys={}, value={}, watermark={}, event_time={}, headers={})",
+            py_repr(py, &self.id)?,
+            py_repr(py, &self.keys)?,
+            bytes_literal(&self.value),
+            self.watermark,
+            self.event_time,
+            py_repr(py, &self.headers)?,
+        ))
     }
 
-    fn __str__(&self) -> String {
-        format!(
-            "Datum(keys={:?}, value={:?}, watermark={}, eventtime={}, id={}, headers={:?})",
-            self.keys,
-            String::from_utf8_lossy(&self.value),
-            self.watermark,
-            self.eventtime,
-            self.id,
-            self.headers
-        )
+    fn __str__(&self, py: Python<'_>) -> PyResult<String> {
+        self.__repr__(py)
     }
 }
 
 impl From<batchmap::Datum> for Datum {
     fn from(value: batchmap::Datum) -> Self {
         Datum::new(
-            value.keys,
-            value.value,
-            value.watermark,
-            value.event_time,
             value.id,
-            value.headers,
+            Some(value.keys),
+            Some(value.value),
+            Some(value.watermark),
+            Some(value.event_time),
+            Some(value.headers),
         )
     }
 }
@@ -175,54 +199,30 @@ impl From<batchmap::Datum> for Datum {
 pub struct BatchResponse {
     #[pyo3(get)]
     pub id: String,
+    #[pyo3(get)]
     pub messages: Vec<Message>,
 }
 
 #[pymethods]
 impl BatchResponse {
+    /// Create a new [BatchResponse] for the given datum id, with optional messages.
     #[new]
-    #[pyo3(signature = (id: "str") -> "BatchResponse")]
-    fn new(id: String) -> Self {
-        Self {
-            id,
-            messages: Vec::new(),
-        }
-    }
-
-    #[staticmethod]
-    #[pyo3(signature = (id: "str") -> "BatchResponse")]
-    fn from_id(id: String) -> Self {
-        Self {
-            id,
-            messages: Vec::new(),
-        }
+    #[pyo3(signature = (id: "str", *messages: "Message") -> "BatchResponse")]
+    fn new(id: String, messages: Vec<Message>) -> Self {
+        Self { id, messages }
     }
 
     #[pyo3(signature = (message))]
     fn append(&mut self, message: Message) {
         self.messages.push(message);
     }
-}
-
-/// A collection of BatchResponse objects for a batch.
-#[pyclass(module = "pynumaflow_lite.batchmapper", from_py_object)]
-#[derive(Clone, Debug)]
-pub struct BatchResponses {
-    pub(crate) responses: Vec<BatchResponse>,
-}
-
-#[pymethods]
-impl BatchResponses {
-    #[new]
-    #[pyo3(signature = () -> "BatchResponses")]
-    fn new() -> Self {
-        Self { responses: vec![] }
-    }
 
-    /// Append a BatchResponse to the collection.
-    #[pyo3(signature = (response: "BatchResponse"))]
-    fn append(&mut self, response: BatchResponse) {
-        self.responses.push(response);
+    fn __repr__(&self, py: Python<'_>) -> PyResult<String> {
+        Ok(format!(
+            "BatchResponse(id={}, messages={})",
+            py_repr(py, &self.id)?,
+            py_repr(py, self.messages.clone())?,
+        ))
     }
 }
 
@@ -271,44 +271,61 @@ impl PyAsyncDatumStream {
 }
 
 /// Async Batch Map Server that can be started from Python code
-#[pyclass(module = "pynumaflow_lite.batchmapper")]
+#[pyclass(name = "_BatchMapAsyncServer", module = "pynumaflow_lite.batchmapper")]
 pub struct BatchMapAsyncServer {
     sock_file: String,
-    info_file: String,
+    server_info_file: String,
     shutdown_tx: Mutex<Option<tokio::sync::oneshot::Sender<()>>>,
 }
 
 #[pymethods]
 impl BatchMapAsyncServer {
     #[new]
-    #[pyo3(signature = (sock_file: "str | None"=batchmap::SOCK_ADDR.to_string(), info_file: "str | None"=batchmap::SERVER_INFO_FILE.to_string()) -> "BatchMapAsyncServer"
-    )]
-    fn new(sock_file: String, info_file: String) -> Self {
+    #[pyo3(signature = (
+        sock_file: "str | None"=None,
+        server_info_file: "str | None"=None,
+    ) -> "_BatchMapAsyncServer")]
+    fn new(sock_file: Option<String>, server_info_file: Option<String>) -> Self {
         Self {
-            sock_file,
-            info_file,
+            sock_file: sock_file.unwrap_or_else(|| batchmap::SOCK_ADDR.to_string()),
+            server_info_file: server_info_file
+                .unwrap_or_else(|| batchmap::SERVER_INFO_FILE.to_string()),
             shutdown_tx: Mutex::new(None),
         }
     }
 
-    #[pyo3(signature = (py_func: "callable") -> "None")]
-    pub fn start<'a>(&self, py: Python<'a>, py_func: Py<PyAny>) -> PyResult<Bound<'a, PyAny>> {
+    #[pyo3(signature = (handler: "callable") -> "None")]
+    pub fn start<'a>(&self, py: Python<'a>, handler: Py<PyAny>) -> PyResult<Bound<'a, PyAny>> {
         let sock_file = self.sock_file.clone();
-        let info_file = self.info_file.clone();
+        let info_file = self.server_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 {
-            // batch server uses the same runner loop and shutdown composition for now
-            // dedicated start is wired below
-            crate::batchmap::server::start(py_func, sock_file, info_file, rx)
-                .await
-                .expect("server failed to start");
-            Ok(())
-        })
+        pyo3_async_runtimes::tokio::future_into_py(
+            py,
+            crate::batchmap::server::start(handler, sock_file, info_file, rx),
+        )
+    }
+
+    /// Wait until the Numaflow IsReady probe succeeds over the batchmap UDS.
+    #[pyo3(signature = (timeout: "float"=30.0) -> "None")]
+    pub fn wait_ready<'a>(&self, py: Python<'a>, timeout: f64) -> PyResult<Bound<'a, PyAny>> {
+        if !timeout.is_finite() || timeout < 0.0 {
+            return Err(pyo3::PyErr::new::<pyo3::exceptions::PyValueError, _>(
+                "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,
+            crate::map::wait_for_ready(sock_file, timeout, "batchmap"),
+        )
     }
 
     #[pyo3(signature = () -> "None")]
@@ -326,7 +343,6 @@ pub(crate) fn populate_py_module(m: &Bound<PyModule>) -> PyResult<()> {
     m.add_class::<Datum>()?;
     m.add_class::<NackOptions>()?;
     m.add_class::<BatchResponse>()?;
-    m.add_class::<BatchResponses>()?;
     m.add_class::<BatchMapAsyncServer>()?;
     m.add_class::<PyAsyncDatumStream>()?;
 
diff --git a/packages/pynumaflow-lite/src/batchmap/server.rs b/packages/pynumaflow-lite/src/batchmap/server.rs
index 25c75210..aca523dd 100644
--- a/packages/pynumaflow-lite/src/batchmap/server.rs
+++ b/packages/pynumaflow-lite/src/batchmap/server.rs
@@ -1,13 +1,31 @@
-// use crate::batchmap::Datum;
+use std::sync::{Arc, Mutex};
+
 use numaflow::batchmap;
 use numaflow::shared::ServerExtras;
-
+use pyo3::exceptions::PyTypeError;
 use pyo3::prelude::*;
-use std::sync::Arc;
+
+use crate::pyrs::{combine_errors, format_error};
 
 pub(crate) struct PyBatchMapRunner {
     pub(crate) event_loop: Arc<Py<PyAny>>,
     pub(crate) py_func: Arc<Py<PyAny>>,
+    pub(crate) errors: Arc<Mutex<Vec<PyErr>>>,
+}
+
+impl PyBatchMapRunner {
+    fn fail(&self, error: PyErr) -> ! {
+        // numaflow calls batchmap() concurrently, so each error belongs to a different
+        // batch. 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}");
+    }
 }
 
 #[tonic::async_trait]
@@ -29,28 +47,53 @@ impl batchmap::BatchMapper for PyBatchMapRunner {
             // When input ends, dropping tx closes the channel
         });
 
-        // Call the Python coroutine: py_func(batch: AsyncIterable[Datum]) -> BatchResponses
-        let fut = Python::attach(|py| {
+        // Call the Python coroutine: py_func(batch: AsyncIterable[Datum]) -> list[BatchResponse]
+        let fut = Python::attach(|py| -> PyResult<_> {
             let locals = pyo3_async_runtimes::TaskLocals::new(self.event_loop.bind(py).clone());
             let py_func = self.py_func.clone();
 
             let stream = crate::batchmap::PyAsyncDatumStream::new_with(rx);
-            let coro = py_func.call1(py, (stream,)).unwrap().into_bound(py);
-            pyo3_async_runtimes::into_future_with_locals(&locals, coro).unwrap()
+            let coro = py_func.call1(py, (stream,))?.into_bound(py);
+            pyo3_async_runtimes::into_future_with_locals(&locals, coro).map_err(|_| {
+                PyErr::new::<PyTypeError, _>(
+                    "batchmap handler must be an async function (coroutine)",
+                )
+            })
         });
 
-        let result = fut.await.unwrap();
+        let fut = match fut {
+            Ok(fut) => fut,
+            Err(error) => self.fail(error),
+        };
+
+        let result = match fut.await {
+            Ok(result) => result,
+            Err(error) => self.fail(error),
+        };
 
         // Ensure forwarder completes
         let _ = forwarder.await;
 
         let responses = Python::attach(|py| {
-            let x: crate::batchmap::BatchResponses = result.extract(py).unwrap();
-            x
+            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(|_| "<unknown>".to_string());
+                PyErr::new::<PyTypeError, _>(format!(
+                    "batchmap handler must return list[BatchResponse], got {type_name}"
+                ))
+            })
         });
 
+        let responses: Vec<crate::batchmap::BatchResponse> = match responses {
+            Ok(responses) => responses,
+            Err(error) => self.fail(error),
+        };
+
         responses
-            .responses
             .into_iter()
             .map(|resp| resp.into())
             .collect::<Vec<batchmap::BatchResponse>>()
@@ -65,14 +108,23 @@ 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 BatchMap UDF. socket={}, server_info={}",
+            sock_file, info_file
+        );
+        move || crate::pyrs::run_asyncio(tx)
+    });
     let event_loop = rx.await.unwrap();
 
+    let errors = Arc::new(Mutex::new(Vec::new()));
+
     let (sig_handle, combined_rx) = crate::pyrs::setup_sig_handler(shutdown_rx);
 
     let py_runner = PyBatchMapRunner {
         py_func: Arc::new(py_func),
         event_loop: event_loop.clone(),
+        errors: Arc::clone(&errors),
     };
 
     let server = numaflow::batchmap::Server::new(py_runner)
@@ -91,11 +143,16 @@ pub(super) async fn start(
         }
     });
 
-    println!("Numaflow Core (batch) has shutdown...");
+    println!("Numaflow BatchMap has shutdown...");
 
     // Wait for the blocking asyncio thread to finish.
     let _ = py_asyncio_loop_handle.await;
 
+    let errors = std::mem::take(&mut *errors.lock().unwrap());
+    if !errors.is_empty() {
+        return Err(Python::attach(|py| combine_errors(py, errors)));
+    }
+
     // if not finished, abort it
     if !sig_handle.is_finished() {
         println!("Aborting signal handler");
diff --git a/packages/pynumaflow-lite/src/lib.rs b/packages/pynumaflow-lite/src/lib.rs
index 6e4f24c5..8fbdcc7e 100644
--- a/packages/pynumaflow-lite/src/lib.rs
+++ b/packages/pynumaflow-lite/src/lib.rs
@@ -108,104 +108,27 @@ fn pynumaflow_lite(py: Python, m: &Bound<PyModule>) -> PyResult<()> {
     m.add_wrapped(pyo3::wrap_pymodule!(sourcetransformer))?;
     m.add_wrapped(pyo3::wrap_pymodule!(sideinputer))?;
 
-    // Ensure it's importable as `pynumaflow_lite.mapper` as well as attribute access
-    let binding = m.getattr("mapper")?;
-    let sub = binding.cast::<PyModule>()?;
-    let fullname = "pynumaflow_lite.mapper";
-    sub.setattr("__name__", fullname)?;
-    py.import("sys")?
-        .getattr("modules")?
-        .set_item(fullname, sub)?;
-
-    // Ensure it's importable as `pynumaflow_lite.batchmapper` as well
-    let binding = m.getattr("batchmapper")?;
-    let sub = binding.cast::<PyModule>()?;
-    let fullname = "pynumaflow_lite.batchmapper";
-    sub.setattr("__name__", fullname)?;
-    py.import("sys")?
-        .getattr("modules")?
-        .set_item(fullname, sub)?;
-
-    // Ensure it's importable as `pynumaflow_lite.mapstreamer` as well
-    let binding = m.getattr("mapstreamer")?;
-    let sub = binding.cast::<PyModule>()?;
-    let fullname = "pynumaflow_lite.mapstreamer";
-    sub.setattr("__name__", fullname)?;
-    py.import("sys")?
-        .getattr("modules")?
-        .set_item(fullname, sub)?;
-
-    // Ensure it's importable as `pynumaflow_lite.reducer` as well
-    let binding = m.getattr("reducer")?;
-    let sub = binding.cast::<PyModule>()?;
-    let fullname = "pynumaflow_lite.reducer";
-    sub.setattr("__name__", fullname)?;
-    py.import("sys")?
-        .getattr("modules")?
-        .set_item(fullname, sub)?;
-
-    // Ensure it's importable as `pynumaflow_lite.session_reducer` as well
-    let binding = m.getattr("session_reducer")?;
-    let sub = binding.cast::<PyModule>()?;
-    let fullname = "pynumaflow_lite.session_reducer";
-    sub.setattr("__name__", fullname)?;
-    py.import("sys")?
-        .getattr("modules")?
-        .set_item(fullname, sub)?;
-
-    // Ensure it's importable as `pynumaflow_lite.reducestreamer` as well
-    let binding = m.getattr("reducestreamer")?;
-    let sub = binding.cast::<PyModule>()?;
-    let fullname = "pynumaflow_lite.reducestreamer";
-    sub.setattr("__name__", fullname)?;
-    py.import("sys")?
-        .getattr("modules")?
-        .set_item(fullname, sub)?;
-
-    // Ensure it's importable as `pynumaflow_lite.accumulator` as well
-    let binding = m.getattr("accumulator")?;
-    let sub = binding.cast::<PyModule>()?;
-    let fullname = "pynumaflow_lite.accumulator";
-    sub.setattr("__name__", fullname)?;
-    py.import("sys")?
-        .getattr("modules")?
-        .set_item(fullname, sub)?;
-
-    // Ensure it's importable as `pynumaflow_lite.sinker` as well
-    let binding = m.getattr("sinker")?;
-    let sub = binding.cast::<PyModule>()?;
-    let fullname = "pynumaflow_lite.sinker";
-    sub.setattr("__name__", fullname)?;
-    py.import("sys")?
-        .getattr("modules")?
-        .set_item(fullname, sub)?;
-
-    // Ensure it's importable as `pynumaflow_lite.sourcer` as well
-    let binding = m.getattr("sourcer")?;
-    let sub = binding.cast::<PyModule>()?;
-    let fullname = "pynumaflow_lite.sourcer";
-    sub.setattr("__name__", fullname)?;
-    py.import("sys")?
-        .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::<PyModule>()?;
-    let fullname = "pynumaflow_lite.sourcetransformer";
-    sub.setattr("__name__", fullname)?;
-    py.import("sys")?
-        .getattr("modules")?
-        .set_item(fullname, sub)?;
-
-    // Ensure it's importable as `pynumaflow_lite.sideinputer` as well
-    let binding = m.getattr("sideinputer")?;
-    let sub = binding.cast::<PyModule>()?;
-    let fullname = "pynumaflow_lite.sideinputer";
-    sub.setattr("__name__", fullname)?;
-    py.import("sys")?
-        .getattr("modules")?
-        .set_item(fullname, sub)?;
+    // Ensure each submodule is importable as `pynumaflow_lite.<name>` as well as attribute access
+    let sys_modules = py.import("sys")?.getattr("modules")?;
+    for name in [
+        "mapper",
+        "batchmapper",
+        "mapstreamer",
+        "reducer",
+        "session_reducer",
+        "reducestreamer",
+        "accumulator",
+        "sinker",
+        "sourcer",
+        "sourcetransformer",
+        "sideinputer",
+    ] {
+        let binding = m.getattr(name)?;
+        let sub = binding.cast::<PyModule>()?;
+        let fullname = format!("pynumaflow_lite.{name}");
+        sub.setattr("__name__", &fullname)?;
+        sys_modules.set_item(&fullname, sub)?;
+    }
 
     Ok(())
 }
diff --git a/packages/pynumaflow-lite/src/map/mod.rs b/packages/pynumaflow-lite/src/map/mod.rs
index 6beaee19..8bcdfdea 100644
--- a/packages/pynumaflow-lite/src/map/mod.rs
+++ b/packages/pynumaflow-lite/src/map/mod.rs
@@ -20,10 +20,7 @@ use pyo3::prelude::*;
 use std::sync::Mutex;
 
 use crate::nack::NackOptions;
-
-pub(crate) fn bytes_literal(value: &[u8]) -> String {
-    format!("b\"{}\"", String::from_utf8_lossy(value).escape_debug())
-}
+use crate::pyrs::bytes_literal;
 
 fn metadata_literal(metadata: &HashMap<String, HashMap<String, Vec<u8>>>) -> String {
     let groups: Vec<String> = metadata
@@ -163,6 +160,10 @@ impl Message {
                 .map_or_else(|| "None".to_string(), metadata_literal),
         )
     }
+
+    fn __str__(&self) -> String {
+        self.__repr__()
+    }
 }
 
 impl From<Message> for map::Message {
@@ -363,10 +364,10 @@ impl MapAsyncServer {
             *guard = Some(tx);
         }
 
-        pyo3_async_runtimes::tokio::future_into_py(py, async move {
-            crate::map::server::start(handler, sock_file, server_info_file, rx).await?;
-            Ok(())
-        })
+        pyo3_async_runtimes::tokio::future_into_py(
+            py,
+            crate::map::server::start(handler, sock_file, server_info_file, rx),
+        )
     }
 
     /// Wait until the Numaflow IsReady probe succeeds over the map UDS.
@@ -381,10 +382,7 @@ impl MapAsyncServer {
         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(())
-        })
+        pyo3_async_runtimes::tokio::future_into_py(py, wait_for_ready(sock_file, timeout, "map"))
     }
 
     /// Trigger server shutdown from Python (idempotent).
diff --git a/packages/pynumaflow-lite/src/map/server.rs b/packages/pynumaflow-lite/src/map/server.rs
index 60d7d83e..9f827494 100644
--- a/packages/pynumaflow-lite/src/map/server.rs
+++ b/packages/pynumaflow-lite/src/map/server.rs
@@ -1,9 +1,11 @@
+use std::sync::{Arc, Mutex};
+
 use numaflow::map;
 use numaflow::shared::ServerExtras;
-
 use pyo3::exceptions::PyTypeError;
 use pyo3::prelude::*;
-use std::sync::{Arc, Mutex};
+
+use crate::pyrs::{combine_errors, format_error};
 
 pub(crate) struct PyMapRunner {
     pub(crate) event_loop: Arc<Py<PyAny>>,
@@ -11,42 +13,6 @@ pub(crate) struct PyMapRunner {
     pub(crate) errors: Arc<Mutex<Vec<PyErr>>>,
 }
 
-// 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(),
-    }
-}
-
-// 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>) -> PyErr {
-    let first = || errors.first().expect("errors is never empty").clone_ref(py);
-
-    if errors.len() == 1 {
-        return first();
-    }
-
-    let Ok(group_type) = py
-        .import("builtins")
-        .and_then(|builtins| builtins.getattr("BaseExceptionGroup"))
-    else {
-        return first();
-    };
-
-    let values: Vec<_> = errors.iter().map(|error| error.value(py).clone()).collect();
-    let message = format!("{} map handler calls failed", values.len());
-
-    match group_type.call1((message, values)) {
-        Ok(group) => PyErr::from_value(group),
-        Err(_) => first(),
-    }
-}
-
 impl PyMapRunner {
     fn fail(&self, error: PyErr) -> ! {
         // numaflow calls map() concurrently, so each error belongs to a different
@@ -137,7 +103,7 @@ pub(super) async fn start(
     let py_map_runner = PyMapRunner {
         py_func: Arc::new(py_func),
         event_loop: event_loop.clone(),
-        errors: errors.clone(),
+        errors: Arc::clone(&errors),
     };
 
     let server = numaflow::map::Server::new(py_map_runner)
diff --git a/packages/pynumaflow-lite/src/nack.rs b/packages/pynumaflow-lite/src/nack.rs
index 95270ad1..3539875e 100644
--- a/packages/pynumaflow-lite/src/nack.rs
+++ b/packages/pynumaflow-lite/src/nack.rs
@@ -41,7 +41,7 @@ impl NackOptions {
         }
     }
 
-    fn __repr__(&self) -> String {
+    pub(crate) fn __repr__(&self) -> String {
         let opt = |v: &Option<String>| v.as_ref().map_or_else(|| "None".to_string(), |s| s.clone());
         format!(
             "NackOptions(delay={}, max_deliveries={}, reason={}, nack_map={:?})",
diff --git a/packages/pynumaflow-lite/src/pyrs.rs b/packages/pynumaflow-lite/src/pyrs.rs
index 23fa999a..f44868ba 100644
--- a/packages/pynumaflow-lite/src/pyrs.rs
+++ b/packages/pynumaflow-lite/src/pyrs.rs
@@ -1,5 +1,7 @@
-use pyo3::{Py, PyAny, Python};
 use std::sync::Arc;
+
+use pyo3::prelude::*;
+use pyo3::{IntoPyObjectExt, Py, PyAny, PyErr, Python};
 use tokio::sync::oneshot::{Receiver, Sender};
 use tokio::task::JoinHandle;
 
@@ -50,3 +52,80 @@ pub(crate) fn setup_sig_handler(shutdown_rx: Receiver<()>) -> (JoinHandle<()>, R
 
     (sig_handle, combined_rx)
 }
+
+// Build the full Python traceback text for the panic message, so the sidecar
+// reports the same failure that Python raises.
+pub(crate) 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(),
+    }
+}
+
+// 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.
+pub(crate) fn combine_errors(py: Python<'_>, errors: Vec<PyErr>) -> PyErr {
+    let first = || errors.first().expect("errors is never empty").clone_ref(py);
+
+    if errors.len() == 1 {
+        return first();
+    }
+
+    let Ok(group_type) = py
+        .import("builtins")
+        .and_then(|builtins| builtins.getattr("BaseExceptionGroup"))
+    else {
+        return first();
+    };
+
+    let values: Vec<_> = errors.iter().map(|error| error.value(py).clone()).collect();
+    let message = format!("{} handler calls failed", values.len());
+
+    match group_type.call1((message, values)) {
+        Ok(group) => PyErr::from_value(group),
+        Err(_) => first(),
+    }
+}
+
+/// Convert a value to a Python object and return its Python `repr()`.
+pub(crate) fn py_repr<'py, T>(py: Python<'py>, value: T) -> PyResult<String>
+where
+    T: IntoPyObject<'py>,
+{
+    Ok(value.into_bound_py_any(py)?.repr()?.to_string())
+}
+
+/// Format bytes the same way as Python `repr(bytes)`.
+pub(crate) fn bytes_literal(value: &[u8]) -> String {
+    use std::fmt::Write;
+
+    let quote = if value.contains(&b'\'') && !value.contains(&b'"') {
+        b'"'
+    } else {
+        b'\''
+    };
+
+    let mut out = String::with_capacity(value.len() + 3);
+    out.push('b');
+    out.push(quote as char);
+    for &byte in value {
+        match byte {
+            b'\\' => out.push_str("\\\\"),
+            b'\t' => out.push_str("\\t"),
+            b'\n' => out.push_str("\\n"),
+            b'\r' => out.push_str("\\r"),
+            _ if byte == quote => {
+                out.push('\\');
+                out.push(quote as char);
+            }
+            0x20..=0x7e => out.push(byte as char),
+            _ => {
+                let _ = write!(out, "\\x{byte:02x}");
+            }
+        }
+    }
+    out.push(quote as char);
+    out
+}
diff --git a/packages/pynumaflow-lite/tests/_test_utils.py b/packages/pynumaflow-lite/tests/_test_utils.py
index fced4816..d86c4f3c 100644
--- a/packages/pynumaflow-lite/tests/_test_utils.py
+++ b/packages/pynumaflow-lite/tests/_test_utils.py
@@ -37,6 +37,34 @@ async def _unused_handler(datums):
     asyncio.run(server.wait_ready(timeout=timeout))
 
 
+def _stop_server(server: subprocess.Popen, timeout: float) -> str:
+    """Stop the server (SIGINT, then SIGKILL after `timeout`) and return its output."""
+    # Request graceful shutdown via SIGINT
+    try:
+        if server.poll() is None:
+            if hasattr(os, "killpg") and server.pid:
+                os.killpg(os.getpgid(server.pid), signal.SIGINT)
+            else:
+                server.send_signal(signal.SIGINT)
+    except Exception:
+        pass
+
+    # Wait for server to exit. communicate() drains the pipe, so a server with a lot of
+    # output cannot block on a full pipe.
+    try:
+        logs, _ = server.communicate(timeout=timeout)
+    except subprocess.TimeoutExpired:
+        try:
+            if hasattr(os, "killpg") and server.pid:
+                os.killpg(os.getpgid(server.pid), signal.SIGKILL)
+            else:
+                server.kill()
+        except Exception:
+            pass
+        logs, _ = server.communicate()
+    return logs or ""
+
+
 def run_python_server_with_rust_client(
     script: str,
     sock_path: Path,
@@ -107,34 +135,16 @@ def run_python_server_with_rust_client(
             env=env,
             timeout=rust_timeout,
         )
-        if rust.returncode != 0:
-            # Dump helpful logs for debugging
-            server_logs = server.stdout.read() if server.stdout else ""
-            pytest.fail(
-                f"Rust client failed: code={rust.returncode}\nStdout:\n{rust.stdout}\nStderr:\n{rust.stderr}\nServer logs so far:\n{server_logs}"
-            )
-
     finally:
-        # Request graceful shutdown via SIGINT
-        try:
-            if server.poll() is None:
-                if hasattr(os, "killpg") and server.pid:
-                    os.killpg(os.getpgid(server.pid), signal.SIGINT)
-                else:
-                    server.send_signal(signal.SIGINT)
-        except Exception:
-            pass
+        # Always stop the server before reading its output. Reading the pipe of a live
+        # server blocks until it exits.
+        server_logs = _stop_server(server, server_shutdown_timeout)
 
-        # Wait for server to exit
-        try:
-            server.wait(timeout=server_shutdown_timeout)
-        except subprocess.TimeoutExpired:
-            try:
-                if hasattr(os, "killpg") and server.pid:
-                    os.killpg(os.getpgid(server.pid), signal.SIGKILL)
-                else:
-                    server.kill()
-            except Exception:
-                pass
+    if rust.returncode != 0:
+        pytest.fail(
+            f"Rust client failed: code={rust.returncode}\nStdout:\n{rust.stdout}\nStderr:\n{rust.stderr}\nServer logs:\n{server_logs}"
+        )
 
-    assert server.returncode == 0, f"Server did not exit cleanly, code={server.returncode}"
+    assert server.returncode == 0, (
+        f"Server did not exit cleanly, code={server.returncode}\nServer logs:\n{server_logs}"
+    )
diff --git a/packages/pynumaflow-lite/tests/examples/batchmap_cat.py b/packages/pynumaflow-lite/tests/examples/batchmap_cat.py
index 64a018e8..5fd343b1 100644
--- a/packages/pynumaflow-lite/tests/examples/batchmap_cat.py
+++ b/packages/pynumaflow-lite/tests/examples/batchmap_cat.py
@@ -1,51 +1,28 @@
 import asyncio
-import collections.abc
-import signal
-from collections.abc import Awaitable, Callable
+from collections.abc import AsyncIterable
 
-from pynumaflow_lite import batchmapper
+from pynumaflow_lite.batchmapper import BatchMapAsyncServer, BatchResponse, Datum, Message
 
 
 async def async_handler(
-    batch: collections.abc.AsyncIterator[batchmapper.Datum],
-) -> batchmapper.BatchResponses:
-    responses = batchmapper.BatchResponses()
+    batch: AsyncIterable[Datum],
+) -> list[BatchResponse]:
+    responses = []
     async for d in batch:
-        resp = batchmapper.BatchResponse.from_id(d.id)
         if d.value == b"bad world":
-            resp.append(batchmapper.Message.message_to_drop())
-            continue
-
-        resp.append(batchmapper.Message(d.value, d.keys))
-        responses.append(resp)
+            responses.append(BatchResponse(d.id, Message.to_drop()))
+        else:
+            responses.append(BatchResponse(d.id, Message(d.value, keys=d.keys)))
     return responses
 
 
-async def start(
-    f: Callable[
-        [collections.abc.AsyncIterator[batchmapper.Datum]],
-        Awaitable[batchmapper.BatchResponses],
-    ],
-):
-    sock_file = "/tmp/var/run/numaflow/batchmap.sock"
-    server_info_file = "/tmp/var/run/numaflow/mapper-server-info"
-    server = batchmapper.BatchMapAsyncServer(sock_file, server_info_file)
-
-    # Register loop-level signal handlers to request graceful shutdown
-    loop = asyncio.get_running_loop()
-    try:
-        loop.add_signal_handler(signal.SIGINT, lambda: server.stop())
-        loop.add_signal_handler(signal.SIGTERM, lambda: server.stop())
-    except (NotImplementedError, RuntimeError):
-        pass
-
-    try:
-        await server.start(f)
-        print("Shutting down gracefully...")
-    except asyncio.CancelledError:
-        server.stop()
-        return
+async def main():
+    await BatchMapAsyncServer(
+        handler=async_handler,
+        sock_file="/tmp/var/run/numaflow/batchmap.sock",
+        server_info_file="/tmp/var/run/numaflow/mapper-server-info",
+    ).serve()
 
 
 if __name__ == "__main__":
-    asyncio.run(start(async_handler))
+    asyncio.run(main())
diff --git a/packages/pynumaflow-lite/tests/examples/batchmap_cat_class.py b/packages/pynumaflow-lite/tests/examples/batchmap_cat_class.py
index f78cfdfe..36dbacf8 100644
--- a/packages/pynumaflow-lite/tests/examples/batchmap_cat_class.py
+++ b/packages/pynumaflow-lite/tests/examples/batchmap_cat_class.py
@@ -1,48 +1,27 @@
 import asyncio
-import signal
-from collections.abc import AsyncIterator, Awaitable, Callable
+from collections.abc import AsyncIterable
 
-from pynumaflow_lite import batchmapper
-from pynumaflow_lite.batchmapper import Message
+from pynumaflow_lite.batchmapper import BatchMapAsyncServer, BatchMapper, BatchResponse, Datum, Message
 
 
-class SimpleBatchCat(batchmapper.BatchMapper):
-    async def handler(self, batch: AsyncIterator[batchmapper.Datum]) -> batchmapper.BatchResponses:
-        responses = batchmapper.BatchResponses()
+class SimpleBatchCat(BatchMapper):
+    async def handler(self, batch: AsyncIterable[Datum]) -> list[BatchResponse]:
+        responses = []
         async for d in batch:
-            resp = batchmapper.BatchResponse(d.id)
             if d.value == b"bad world":
-                resp.append(Message.message_to_drop())
-                continue
-
-            resp.append(Message(d.value, d.keys))
-            responses.append(resp)
+                responses.append(BatchResponse(d.id, Message.to_drop()))
+            else:
+                responses.append(BatchResponse(d.id, Message(d.value, d.keys)))
         return responses
 
 
-async def start(
-    f: Callable[[AsyncIterator[batchmapper.Datum]], Awaitable[batchmapper.BatchResponses]],
-):
-    sock_file = "/tmp/var/run/numaflow/batchmap.sock"
-    server_info_file = "/tmp/var/run/numaflow/mapper-server-info"
-    server = batchmapper.BatchMapAsyncServer(sock_file, server_info_file)
-
-    # Register loop-level signal handlers so we control shutdown and avoid asyncio.run
-    loop = asyncio.get_running_loop()
-    try:
-        loop.add_signal_handler(signal.SIGINT, lambda: server.stop())
-        loop.add_signal_handler(signal.SIGTERM, lambda: server.stop())
-    except (NotImplementedError, RuntimeError):
-        pass
-
-    try:
-        await server.start(f)
-        print("Shutting down gracefully...")
-    except asyncio.CancelledError:
-        server.stop()
-        return
+async def main():
+    await BatchMapAsyncServer(
+        SimpleBatchCat(),
+        sock_file="/tmp/var/run/numaflow/batchmap.sock",
+        server_info_file="/tmp/var/run/numaflow/mapper-server-info",
+    ).serve()
 
 
 if __name__ == "__main__":
-    async_handler = SimpleBatchCat()
-    asyncio.run(start(async_handler))
+    asyncio.run(main())