[nf-concurrent-batch-proc] feat: adding a minimal example for BatchMapper with BatchMapAsyncServer

.dockerignore

@@ -0,0 +1,6 @@
+.venv
+*.yaml
+.git
+__pycache__
+*.pyc
+.DS_Store

packages/pynumaflow/examples/batchmap/concurrent_sink/Dockerfile

@@ -0,0 +1,38 @@
+####################################################################################################
+# Stage 1: Builder - installs all dependencies using uv
+####################################################################################################
+FROM ghcr.io/astral-sh/uv:python3.13-trixie AS builder
+
+ENV PYSETUP_PATH="/opt/pysetup"
+WORKDIR $PYSETUP_PATH
+
+COPY pyproject.toml uv.lock README.md ./
+COPY pynumaflow/ ./pynumaflow/
+
+ENV EXAMPLE_PATH="$PYSETUP_PATH/examples/batchmap/concurrent_sink"
+COPY examples/batchmap/concurrent_sink/ $EXAMPLE_PATH/
+
+WORKDIR $EXAMPLE_PATH
+RUN uv sync --no-dev --no-install-project --frozen
+
+####################################################################################################
+# Stage 2: Runtime - clean image with only installed packages
+####################################################################################################
+FROM ghcr.io/astral-sh/uv:python3.13-trixie AS udf
+
+ENV PYSETUP_PATH="/opt/pysetup"
+ENV EXAMPLE_PATH="$PYSETUP_PATH/examples/batchmap/concurrent_sink"
+
+WORKDIR $EXAMPLE_PATH
+COPY --from=builder $EXAMPLE_PATH/.venv $EXAMPLE_PATH/.venv
+COPY --from=builder $EXAMPLE_PATH/ $EXAMPLE_PATH/
+
+# NOTE: We cannot use "uv run python example.py" here because uv run reads the
+# example's pyproject.toml, finds the pynumaflow path source (path = "../../../"),
+# and tries to resolve it. In the runtime stage, the parent pynumaflow source tree
+# is not present (by design, to keep the image small), so uv run fails.
+# Instead, we activate the pre-built .venv via PATH and run python directly.
+ENV PATH="$EXAMPLE_PATH/.venv/bin:$PATH"
+CMD ["python", "example.py"]
+
+EXPOSE 5000

packages/pynumaflow/examples/batchmap/concurrent_sink/Makefile

@@ -0,0 +1,22 @@
+TAG ?= stable
+PUSH ?= false
+IMAGE_REGISTRY = quay.io/numaio/numaflow-python/batch-map-concurrent-sink:${TAG}
+DOCKER_FILE_PATH = examples/batchmap/concurrent_sink/Dockerfile
+
+.PHONY: update
+update:
+	uv lock --check || uv lock
+
+.PHONY: image-push
+image-push: update
+	cd ../../../ && docker buildx build \
+	-f ${DOCKER_FILE_PATH} \
+	-t ${IMAGE_REGISTRY} \
+	--platform linux/amd64,linux/arm64 . --push
+
+.PHONY: image
+image: update
+	cd ../../../ && podman build --tls-verify=false \
+	-f ${DOCKER_FILE_PATH} \
+	-t ${IMAGE_REGISTRY} .
+	@if [ "$(PUSH)" = "true" ]; then docker push ${IMAGE_REGISTRY}; fi

packages/pynumaflow/examples/batchmap/concurrent_sink/README.md

@@ -0,0 +1,92 @@
+# Concurrent Sink Example
+
+Demonstrates concurrent processing of batch messages using `ProcessPoolExecutor` with `BatchMapAsyncServer`.
+Use batch processing only when it makes sense. In some scenarios, batch processing may **not** be the most efficient approach, and processing data items one by one could be a better option. 
+The burden of concurrent processing of the data will rely on the UDF implementation in this use case.
+
+## Overview
+
+This example shows a workload by offloading CPU-intensive work to a pool of worker processes. Each message in a batch is processed independently and in parallel, improving throughput for compute-heavy operations.
+
+## Key Features
+
+- **Parallel Processing**: Uses `ProcessPoolExecutor` to process multiple messages concurrently
+- **Async/Await**: Leverages `asyncio` for non-blocking task submission and collection
+- **Cross-Process Safety**: Returns only primitive types from worker processes (strings, bytes, booleans)
+- **Graceful Degradation**: Drops messages on processing errors
+- **Configurable Workers**: Adjust worker count via `NUM_WORKERS` environment variable
+
+## Architecture
+
+1. **Main Process (AsyncIO Loop)**
+   - Receives batch of datums from Numaflow
+   - Submits each message to `ProcessPoolExecutor` via `loop.run_in_executor()`
+   - Collects results concurrently with `asyncio.gather()`
+   - Constructs and returns `BatchResponses`
+
+2. **Worker Processes**
+   - Execute `_process_message_task()` in isolation
+   - Perform CPU-intensive work (hash computation in this example)
+   - Return primitive types only (no Pynumaflow objects)
+
+## Usage
+
+```bash
+# Start with default worker count (CPU count)
+python example.py
+
+# Start with custom worker count
+NUM_WORKERS=4 python example.py
+```
+
+## Message Flow
+
+```
+Input Batch → Async Drain → Submit to Workers → Collect Results → BatchResponses
+                ↓                   ↓                    ↓                 ↓
+        [D1, D2, D3, ...]  Worker Pool (P1-Pn)  Task Futures  [BR1, BR2, BR3, ...]
+```
+
+## Processing Example
+
+**Input Message**: `"hello world"`
+
+**Processing**:
+1. Decode UTF-8
+2. Compute SHA256 hash 100 times iteratively
+3. Return processed result with truncated hash
+
+**Output**: `"PROCESSED[hello world]-HASH[<truncated_hash>]"`
+
+## Error Handling
+
+- **Process Timeout**: Drops message (returns empty payload)
+- **Exception in Worker**: Logs error, drops message
+- **Invalid Input**: Gracefully handled with try-except
+
+## Configuration
+
+### Environment Variables
+
+- `NUM_WORKERS`: Number of worker processes (default: CPU count)
+
+### Resource Tuning
+
+For CPU-bound work:
+- Set `NUM_WORKERS` ≈ CPU cores
+- Larger batches improve throughput
+- Monitor worker process memory usage
+
+## Comparison to Alternatives
+
+| Approach | Pros | Cons |
+|----------|------|------|
+| **ProcessPool** (this) | True parallelism, CPU-bound work | Process overhead, IPC serialization |
+| **ThreadPool** | Lower overhead | GIL contention, not for CPU work |
+| **AsyncIO** | Lightweight, I/O-bound | Single-threaded, no true parallelism |
+
+## Related Examples
+
+- `examples/map/multiproc_map/`: Single-message multiprocessing mapper
+- `examples/batchmap/flatmap/`: Basic batch mapping without parallelism
+- `examples/sink/async_log/`: Async sink without process pool

packages/pynumaflow/examples/batchmap/concurrent_sink/example.py

@@ -0,0 +1,147 @@
+import asyncio
+import hashlib
+import logging
+import os
+from collections.abc import AsyncIterable
+from concurrent.futures import ProcessPoolExecutor
+from typing import Tuple
+
+from pynumaflow.batchmapper import (
+    BatchMapper,
+    BatchResponse,
+    BatchResponses,
+    Datum,
+)
+from pynumaflow.batchmapper import BatchMapAsyncServer
+from pynumaflow.mapper import Message
+
+logging.basicConfig(level=logging.INFO)
+_LOGGER = logging.getLogger(__name__)
+
+# Process-level executor
+_EXECUTOR = None
+
+
+def get_executor(max_workers: int = None) -> ProcessPoolExecutor:
+    """Get or create global ProcessPoolExecutor."""
+    global _EXECUTOR
+    if _EXECUTOR is None:
+        if max_workers is None:
+            max_workers = os.cpu_count() or 2
+        _EXECUTOR = ProcessPoolExecutor(max_workers=max_workers)
+    return _EXECUTOR
+
+
+def _process_message_task(task_data: Tuple[str, list, bytes]) -> Tuple[str, list, bytes, bool]:
+    """Worker function executed in ProcessPoolExecutor child process.
+
+    Returns primitive types (msg_id, keys, payload_bytes, should_drop) to safely
+    cross process boundaries.
+    """
+    msg_id, keys, datum_bytes = task_data
+    pid = os.getpid()
+
+    try:
+        # Decode message
+        message_str = (
+            datum_bytes.decode("utf-8")
+            if isinstance(datum_bytes, bytes)
+            else datum_bytes
+        )
+        _LOGGER.info(f"[Worker PID: {pid}] Processing message: {message_str}")
+
+        # CPU-intensive operation: compute hash multiple times
+        processed = message_str
+        for _ in range(128):
+            processed = hashlib.sha256(processed.encode()).hexdigest()
+
+        # Simulate transformation
+        result = f"PROCESSED[{message_str}]-HASH[{processed[:16]}]"
+        result_bytes = result.encode("utf-8")
+
+        _LOGGER.info(f"[Worker PID: {pid}] Processed msg_id={msg_id}: {result}")
+        return msg_id, keys, result_bytes, False
+
+    except Exception as e:
+        _LOGGER.error(f"[Worker PID: {pid}] Error processing msg_id={msg_id}: {e}")
+        return msg_id, keys, b"", True
+
+
+class ConcurrentSink(BatchMapper):
+    """BatchMapper that processes messages concurrently using ProcessPoolExecutor."""
+
+    def __init__(self, max_workers: int = None):
+        """Initialize sink with process pool.
+
+        Args:
+            max_workers: Number of worker processes. Defaults to CPU count.
+        """
+        self.executor = get_executor(max_workers)
+        _LOGGER.info(f"Initialized ConcurrentSink with ProcessPoolExecutor (PID={os.getpid()})")
+
+    async def handler(self, datums: AsyncIterable[Datum]) -> BatchResponses:
+        """Process batch of datums concurrently in worker processes.
+
+        Args:
+            datums: AsyncIterable of Datum objects from the stream.
+
+        Returns:
+            BatchResponses with results from all processed messages.
+        """
+        responses = BatchResponses()
+        loop = asyncio.get_running_loop()
+
+        # Collect all datums from the async iterable
+        datums_list = []
+        async for datum in datums:
+            datums_list.append(datum)
+
+        _LOGGER.info(f"Processing batch of {len(datums_list)} messages")
+
+        if not datums_list:
+            return responses
+
+        # Submit all tasks to process pool concurrently
+        async def submit_task(datum: Datum):
+            """Submit a single message to worker pool."""
+            task_data = (datum.id, datum.keys(), datum.value)
+            return await loop.run_in_executor(self.executor, _process_message_task, task_data)
+
+        # Wait for all results
+        results = await asyncio.gather(
+            *[submit_task(d) for d in datums_list],
+            return_exceptions=True,
+        )
+
+        # Construct responses from results
+        for datum, res in zip(datums_list, results):
+            batch_resp = BatchResponse.from_id(datum.id)
+
+            if isinstance(res, Exception):
+                _LOGGER.error(f"Task failed for datum ID: {datum.id}, error: {res}")
+                batch_resp.append(Message.to_drop())
+            else:
+                msg_id, keys, payload_bytes, should_drop = res
+                if should_drop:
+                    batch_resp.append(Message.to_drop())
+                else:
+                    batch_resp.append(Message(value=payload_bytes, keys=keys))
+
+            responses.append(batch_resp)
+
+        _LOGGER.info(f"Completed processing, generated {len(responses)} responses")
+        return responses
+
+
+if __name__ == "__main__":
+    """
+    Example of starting a concurrent sink with process pool.
+
+    To configure the number of worker processes, set NUM_WORKERS env var:
+        NUM_WORKERS=4 python example.py
+    NB: NUM_WORKERS is a user created ENV variable, not a platform ENV varible
+    """
+    max_workers = int(os.getenv("NUM_WORKERS", str(os.cpu_count() or 2)))
+    sink = ConcurrentSink(max_workers=max_workers)
+    grpc_server = BatchMapAsyncServer(sink)
+    grpc_server.start()

packages/pynumaflow/examples/batchmap/concurrent_sink/pipeline.yaml

@@ -0,0 +1,33 @@
+apiVersion: numaflow.numaproj.io/v1alpha1
+kind: Pipeline
+metadata:
+  name: concurrent-sink
+spec:
+  vertices:
+    - name: in
+      source:
+        # A self data generating source
+        generator:
+          rpu: 500
+          duration: 1s
+    - name: batch-concurrent-sink
+      partitions: 2
+      scale:
+        min: 1
+      udf:
+        container:
+          image: quay.io/numaio/numaflow-python/batch-map-concurrent-sink:stable
+          imagePullPolicy: Always
+          env:
+            - name: NUM_WORKERS
+              value: "2"
+    - name: sink
+      scale:
+        min: 1
+      sink:
+        log: {}
+  edges:
+    - from: in
+      to: batch-concurrent-sink
+    - from: batch-concurrent-sink
+      to: sink

packages/pynumaflow/examples/batchmap/concurrent_sink/pyproject.toml

@@ -0,0 +1,18 @@
+[project]
+name = "batch-map-concurrent-sink"
+version = "0.2.4"
+description = ""
+requires-python = ">=3.13"
+dependencies = [
+    "pynumaflow",
+]
+
+[tool.uv.sources]
+pynumaflow = { path = "../../../" }
+
+[tool.hatch.build.targets.wheel]
+packages = []
+
+[build-system]
+requires = ["hatchling"]
+build-backend = "hatchling.build"