feat: end to end wiring for per message nack (#365)

Signed-off-by: Vaibhav Tiwari <vaibhav.tiwari33@gmail.com>
Signed-off-by: Abdullah Yildirim <a3hadi@protonmail.com>
Signed-off-by: GitHub <noreply@github.com>
Co-authored-by: Vaibhav Kant Tiwari <vaibhav.tiwari33@gmail.com>
Co-authored-by: github-actions[bot] <41898282+github-actions[bot]@users.noreply.github.com>
Co-authored-by: vaibhavtiwari33 <vaibhavtiwari33@users.noreply.github.com>

Author: 3 people

packages/pynumaflow/pynumaflow/_constants.py

@@ -54,6 +54,7 @@
 STREAM_EOF = "EOF"
 DELIMITER = ":"
 DROP = "U+005C__DROP__"
+NACK = "U+005C__NACK__"
 
 _PROCESS_COUNT = os.cpu_count()
 # Cap max value to 16

packages/pynumaflow/pynumaflow/_nack.py

@@ -0,0 +1,52 @@
+from __future__ import annotations
+
+from dataclasses import dataclass, field
+
+from pynumaflow.proto.common import nack_options_pb2
+
+
+@dataclass
+class NackOptions:
+    """Per-message redelivery options for a nack.
+
+    Args:
+        delay: the redelivery delay in milliseconds.
+        max_deliveries: the maximum number of redelivery attempts.
+        reason: a human-readable reason for nacking the message.
+        nack_map: a generic map of string key-value pairs used to pass
+            properties/configurations for nacking back to the source (e.g. for
+            sources like SQS, Pulsar or JetStream that support NACK).
+    """
+
+    delay: int | None = None
+    max_deliveries: int | None = None
+    reason: str | None = None
+    nack_map: dict[str, str] = field(default_factory=dict)
+
+    def _to_proto(self) -> nack_options_pb2.NackOptions:
+        return nack_options_pb2.NackOptions(
+            reason=self.reason,
+            max_deliveries=self.max_deliveries,
+            delay=self.delay,
+            nack_map=self.nack_map,
+        )
+
+
+def _nack_options_to_proto(
+    opts: NackOptions | None,
+) -> nack_options_pb2.NackOptions | None:
+    if opts is None:
+        return None
+    return opts._to_proto()
+
+
+def _nack_options_from_proto(
+    proto: nack_options_pb2.NackOptions,
+) -> NackOptions:
+    return NackOptions(
+        delay=proto.delay if proto.HasField("delay") else None,
+        max_deliveries=proto.max_deliveries if proto.HasField("max_deliveries") else None,
+        reason=proto.reason if proto.HasField("reason") else None,
+        # proto3 map fields do not support HasField; an unset map is simply empty.
+        nack_map=dict(proto.nack_map),
+    )

packages/pynumaflow/pynumaflow/batchmapper/__init__.py

@@ -8,6 +8,7 @@
     BatchResponse,
 )
 from pynumaflow.batchmapper.async_server import BatchMapAsyncServer
+from pynumaflow._nack import NackOptions
 
 __all__ = [
     "Message",
@@ -17,4 +18,5 @@
     "BatchMapper",
     "BatchResponses",
     "BatchResponse",
+    "NackOptions",
 ]

packages/pynumaflow/pynumaflow/batchmapper/_dtypes.py

@@ -5,7 +5,8 @@
 from typing import TypeAlias, TypeVar
 from collections.abc import AsyncIterable, Callable
 
-from pynumaflow._constants import DROP
+from pynumaflow._constants import DROP, NACK
+from pynumaflow._nack import NackOptions
 from pynumaflow._validate import _validate_message_fields
 
 M = TypeVar("M", bound="Message")
@@ -27,6 +28,7 @@ class Message:
     _value: bytes
     _keys: list[str]
     _tags: list[str]
+    _nack_options: NackOptions | None
 
     def __init__(self, value: bytes, keys: list[str] | None = None, tags: list[str] | None = None):
         """
@@ -36,12 +38,23 @@ def __init__(self, value: bytes, keys: list[str] | None = None, tags: list[str]
         self._keys = keys or []
         self._tags = tags or []
         self._value = value or b""
+        self._nack_options = None
 
     # returns the Message Object which will be dropped
     @classmethod
     def to_drop(cls: type[M]) -> M:
         return cls(b"", None, [DROP])
 
+    @classmethod
+    def to_nack(cls: type[M], opts: NackOptions | None = None) -> M:
+        m = cls(b"", None, [NACK])
+        m._nack_options = opts
+        return m
+
+    @property
+    def nack_options(self) -> NackOptions | None:
+        return self._nack_options
+
     @property
     def value(self) -> bytes:
         return self._value

packages/pynumaflow/pynumaflow/batchmapper/servicer/async_servicer.py

@@ -10,6 +10,7 @@
 from pynumaflow.shared.server import update_context_err
 from pynumaflow.types import NumaflowServicerContext
 from pynumaflow._constants import _LOGGER, STREAM_EOF, ERR_UDF_EXCEPTION_STRING
+from pynumaflow._nack import _nack_options_to_proto
 
 
 class AsyncBatchMapServicer(map_pb2_grpc.MapServicer):
@@ -80,7 +81,10 @@ async def MapFn(
                         for msg in batch_response.messages:
                             single_req_resp.append(
                                 map_pb2.MapResponse.Result(
-                                    keys=msg.keys, value=msg.value, tags=msg.tags
+                                    keys=msg.keys,
+                                    value=msg.value,
+                                    tags=msg.tags,
+                                    nack_options=_nack_options_to_proto(msg.nack_options),
                                 )
                             )
                         # send the response for a given ID back to the stream

packages/pynumaflow/pynumaflow/mapper/__init__.py

@@ -4,6 +4,7 @@
 
 from pynumaflow.mapper._dtypes import Message, Messages, Datum, DROP, Mapper
 from pynumaflow._metadata import UserMetadata, SystemMetadata
+from pynumaflow._nack import NackOptions
 
 __all__ = [
     "Message",
@@ -16,4 +17,5 @@
     "MapMultiprocServer",
     "UserMetadata",
     "SystemMetadata",
+    "NackOptions",
 ]

packages/pynumaflow/pynumaflow/mapper/_dtypes.py

@@ -6,7 +6,8 @@
 from collections.abc import Callable
 from warnings import warn
 
-from pynumaflow._constants import DROP
+from pynumaflow._constants import DROP, NACK
+from pynumaflow._nack import NackOptions
 from pynumaflow._metadata import UserMetadata, SystemMetadata
 from pynumaflow._validate import _validate_message_fields
 
@@ -30,6 +31,7 @@ class Message:
     _keys: list[str]
     _tags: list[str]
     _user_metadata: UserMetadata
+    _nack_options: NackOptions | None
 
     def __init__(
         self,
@@ -46,12 +48,23 @@ def __init__(
         self._tags = tags or []
         self._value = value or b""
         self._user_metadata = user_metadata or UserMetadata()
+        self._nack_options = None
 
     # returns the Message Object which will be dropped
     @classmethod
     def to_drop(cls: type[M]) -> M:
         return cls(b"", None, [DROP])
 
+    @classmethod
+    def to_nack(cls: type[M], opts: NackOptions | None = None) -> M:
+        m = cls(b"", None, [NACK])
+        m._nack_options = opts
+        return m
+
+    @property
+    def nack_options(self) -> NackOptions | None:
+        return self._nack_options
+
     @property
     def value(self) -> bytes:
         return self._value

packages/pynumaflow/pynumaflow/mapper/_servicer/_async_servicer.py

@@ -7,6 +7,7 @@
 from pynumaflow._constants import _LOGGER, STREAM_EOF, ERR_UDF_EXCEPTION_STRING
 from pynumaflow.mapper._dtypes import MapAsyncCallable, Datum, MapError, Message, Messages
 from pynumaflow._metadata import _user_and_system_metadata_from_proto
+from pynumaflow._nack import _nack_options_to_proto
 from pynumaflow.proto.mapper import map_pb2, map_pb2_grpc
 from pynumaflow.shared.server import update_context_err
 from pynumaflow.types import NumaflowServicerContext
@@ -143,6 +144,7 @@ async def _invoke_map(self, req: map_pb2.MapRequest, result_queue: NonBlockingIt
                         value=msg.value,
                         tags=msg.tags,
                         metadata=msg.user_metadata._to_proto(),
+                        nack_options=_nack_options_to_proto(msg.nack_options),
                     )
                 )
             await result_queue.put(map_pb2.MapResponse(results=datums, id=req.id))

packages/pynumaflow/pynumaflow/mapper/_servicer/_sync_servicer.py

@@ -9,6 +9,7 @@
 
 from pynumaflow._constants import NUM_THREADS_DEFAULT, STREAM_EOF, _LOGGER, ERR_UDF_EXCEPTION_STRING
 from pynumaflow.mapper._dtypes import MapSyncCallable, Datum, MapError
+from pynumaflow._nack import _nack_options_to_proto
 from pynumaflow.proto.mapper import map_pb2, map_pb2_grpc
 from pynumaflow.shared.synciter import SyncIterator
 from pynumaflow.types import NumaflowServicerContext
@@ -155,6 +156,7 @@ def _invoke_map(
                         value=resp.value,
                         tags=resp.tags,
                         metadata=resp.user_metadata._to_proto(),
+                        nack_options=_nack_options_to_proto(resp.nack_options),
                     )
                 )
             result_queue.put(map_pb2.MapResponse(results=results, id=request.id))

packages/pynumaflow/pynumaflow/mapstreamer/__init__.py

@@ -2,6 +2,7 @@
 
 from pynumaflow.mapstreamer._dtypes import Message, Messages, Datum, MapStreamer
 from pynumaflow.mapstreamer.async_server import MapStreamAsyncServer
+from pynumaflow._nack import NackOptions
 
 __all__ = [
     "Message",
@@ -10,4 +11,5 @@
     "DROP",
     "MapStreamAsyncServer",
     "MapStreamer",
+    "NackOptions",
 ]