chore: ud-container env variable for formalising the errors (numaproj#220)

Author: adarsh0728

pynumaflow/_constants.py

@@ -5,16 +5,12 @@
 from pynumaflow import setup_logging
 
 SIDE_INPUT_DIR_PATH = "/var/numaflow/side-inputs"
+ENV_UD_CONTAINER_TYPE = "NUMAFLOW_UD_CONTAINER_TYPE"
 
-# UDF execution error prefixes
-ERR_SOURCE_EXCEPTION = "UDF_EXECUTION_ERROR(source)"
-ERR_TRANSFORMER_EXCEPTION = "UDF_EXECUTION_ERROR(transformer)"
-ERR_SINK_EXCEPTION = "UDF_EXECUTION_ERROR(sink)"
-ERR_MAP_STREAM_EXCEPTION = "UDF_EXECUTION_ERROR(mapstream)"
-ERR_MAP_EXCEPTION = "UDF_EXECUTION_ERROR(map)"
-ERR_BATCH_MAP_EXCEPTION = "UDF_EXECUTION_ERROR(batchmap)"
-ERR_REDUCE_EXCEPTION = "UDF_EXECUTION_ERROR(reduce)"
-ERR_SIDE_INPUT_RETRIEVAL_EXCEPTION = "UDF_EXECUTION_ERROR(sideinput)"
+# Get container type from env var, default to unknown-container
+CONTAINER_TYPE = os.getenv(ENV_UD_CONTAINER_TYPE, "unknown-container")
+# UDF exception error string with container type
+ERR_UDF_EXCEPTION_STRING = f"UDF_EXECUTION_ERROR({CONTAINER_TYPE})"
 
 # Socket configs
 MAP_SOCK_PATH = "/var/run/numaflow/map.sock"

pynumaflow/batchmapper/servicer/async_servicer.py

@@ -9,7 +9,7 @@
 from pynumaflow.shared.asynciter import NonBlockingIterator
 from pynumaflow.shared.server import handle_async_error
 from pynumaflow.types import NumaflowServicerContext
-from pynumaflow._constants import _LOGGER, STREAM_EOF, ERR_BATCH_MAP_EXCEPTION
+from pynumaflow._constants import _LOGGER, STREAM_EOF, ERR_UDF_EXCEPTION_STRING
 
 
 class AsyncBatchMapServicer(map_pb2_grpc.MapServicer):
@@ -98,7 +98,7 @@ async def MapFn(
 
         except BaseException as err:
             _LOGGER.critical("UDFError, re-raising the error", exc_info=True)
-            await handle_async_error(context, err, ERR_BATCH_MAP_EXCEPTION)
+            await handle_async_error(context, err, ERR_UDF_EXCEPTION_STRING)
             return
 
     async def IsReady(

pynumaflow/mapper/_servicer/_async_servicer.py

@@ -4,7 +4,7 @@
 from google.protobuf import empty_pb2 as _empty_pb2
 from pynumaflow.shared.asynciter import NonBlockingIterator
 
-from pynumaflow._constants import _LOGGER, STREAM_EOF, ERR_MAP_EXCEPTION
+from pynumaflow._constants import _LOGGER, STREAM_EOF, ERR_UDF_EXCEPTION_STRING
 from pynumaflow.mapper._dtypes import MapAsyncCallable, Datum, MapError
 from pynumaflow.proto.mapper import map_pb2, map_pb2_grpc
 from pynumaflow.shared.server import handle_async_error
@@ -56,7 +56,7 @@ async def MapFn(
             async for msg in consumer:
                 # If the message is an exception, we raise the exception
                 if isinstance(msg, BaseException):
-                    await handle_async_error(context, msg, ERR_MAP_EXCEPTION)
+                    await handle_async_error(context, msg, ERR_UDF_EXCEPTION_STRING)
                     return
                 # Send window response back to the client
                 else:
@@ -65,7 +65,7 @@ async def MapFn(
             await producer
         except BaseException as e:
             _LOGGER.critical("UDFError, re-raising the error", exc_info=True)
-            await handle_async_error(context, e, ERR_MAP_EXCEPTION)
+            await handle_async_error(context, e, ERR_UDF_EXCEPTION_STRING)
             return
 
     async def _process_inputs(

pynumaflow/mapper/_servicer/_sync_servicer.py

@@ -5,7 +5,7 @@
 from google.protobuf import empty_pb2 as _empty_pb2
 from pynumaflow.shared.server import exit_on_error
 
-from pynumaflow._constants import NUM_THREADS_DEFAULT, STREAM_EOF, _LOGGER, ERR_MAP_EXCEPTION
+from pynumaflow._constants import NUM_THREADS_DEFAULT, STREAM_EOF, _LOGGER, ERR_UDF_EXCEPTION_STRING
 from pynumaflow.mapper._dtypes import MapSyncCallable, Datum, MapError
 from pynumaflow.proto.mapper import map_pb2, map_pb2_grpc
 from pynumaflow.shared.synciter import SyncIterator
@@ -58,7 +58,7 @@ def MapFn(
                 if isinstance(res, BaseException):
                     # Terminate the current server process due to exception
                     exit_on_error(
-                        context, f"{ERR_MAP_EXCEPTION}: {repr(res)}", parent=self.multiproc
+                        context, f"{ERR_UDF_EXCEPTION_STRING}: {repr(res)}", parent=self.multiproc
                     )
                     return
                 # return the result
@@ -71,7 +71,9 @@ def MapFn(
         except BaseException as err:
             _LOGGER.critical("UDFError, re-raising the error", exc_info=True)
             # Terminate the current server process due to exception
-            exit_on_error(context, f"{ERR_MAP_EXCEPTION}: {repr(err)}", parent=self.multiproc)
+            exit_on_error(
+                context, f"{ERR_UDF_EXCEPTION_STRING}: {repr(err)}", parent=self.multiproc
+            )
             return
 
     def _process_requests(

pynumaflow/mapstreamer/servicer/async_servicer.py

@@ -7,7 +7,7 @@
 from pynumaflow.proto.mapper import map_pb2_grpc, map_pb2
 from pynumaflow.shared.server import handle_async_error
 from pynumaflow.types import NumaflowServicerContext
-from pynumaflow._constants import _LOGGER, ERR_MAP_STREAM_EXCEPTION
+from pynumaflow._constants import _LOGGER, ERR_UDF_EXCEPTION_STRING
 
 
 class AsyncMapStreamServicer(map_pb2_grpc.MapServicer):
@@ -59,7 +59,7 @@ async def MapFn(
                 yield map_pb2.MapResponse(status=map_pb2.TransmissionStatus(eot=True), id=req.id)
         except BaseException as err:
             _LOGGER.critical("UDFError, re-raising the error", exc_info=True)
-            await handle_async_error(context, err, ERR_MAP_STREAM_EXCEPTION)
+            await handle_async_error(context, err, ERR_UDF_EXCEPTION_STRING)
             return
 
     async def __invoke_map_stream(self, keys: list[str], req: Datum):

pynumaflow/reducer/servicer/async_servicer.py

@@ -3,7 +3,7 @@
 
 from google.protobuf import empty_pb2 as _empty_pb2
 
-from pynumaflow._constants import _LOGGER, ERR_REDUCE_EXCEPTION
+from pynumaflow._constants import _LOGGER, ERR_UDF_EXCEPTION_STRING
 from pynumaflow.proto.reducer import reduce_pb2, reduce_pb2_grpc
 from pynumaflow.reducer._dtypes import (
     Datum,
@@ -105,7 +105,7 @@ async def ReduceFn(
             _LOGGER.critical("Reduce Error", exc_info=True)
             # Send a context abort signal for the rpc, this is required for numa container to get
             # the correct grpc error
-            await handle_async_error(context, e, ERR_REDUCE_EXCEPTION)
+            await handle_async_error(context, e, ERR_UDF_EXCEPTION_STRING)
 
         # send EOF to all the tasks once the request iterator is exhausted
         # This will signal the tasks to stop reading the data on their
@@ -136,7 +136,7 @@ async def ReduceFn(
             _LOGGER.critical("Reduce Error", exc_info=True)
             # Send a context abort signal for the rpc, this is required for numa container to get
             # the correct grpc error
-            await handle_async_error(context, e, ERR_REDUCE_EXCEPTION)
+            await handle_async_error(context, e, ERR_UDF_EXCEPTION_STRING)
 
     async def IsReady(
         self, request: _empty_pb2.Empty, context: NumaflowServicerContext

pynumaflow/reducestreamer/servicer/async_servicer.py

@@ -4,7 +4,7 @@
 
 from google.protobuf import empty_pb2 as _empty_pb2
 
-from pynumaflow._constants import ERR_REDUCE_EXCEPTION
+from pynumaflow._constants import ERR_UDF_EXCEPTION_STRING
 from pynumaflow.proto.reducer import reduce_pb2, reduce_pb2_grpc
 from pynumaflow.reducestreamer._dtypes import (
     Datum,
@@ -95,20 +95,20 @@ async def ReduceFn(
             async for msg in consumer:
                 # If the message is an exception, we raise the exception
                 if isinstance(msg, BaseException):
-                    await handle_async_error(context, msg, ERR_REDUCE_EXCEPTION)
+                    await handle_async_error(context, msg, ERR_UDF_EXCEPTION_STRING)
                     return
                 # Send window EOF response or Window result response
                 # back to the client
                 else:
                     yield msg
         except BaseException as e:
-            await handle_async_error(context, e, ERR_REDUCE_EXCEPTION)
+            await handle_async_error(context, e, ERR_UDF_EXCEPTION_STRING)
             return
         # Wait for the process_input_stream task to finish for a clean exit
         try:
             await producer
         except BaseException as e:
-            await handle_async_error(context, e, ERR_REDUCE_EXCEPTION)
+            await handle_async_error(context, e, ERR_UDF_EXCEPTION_STRING)
             return
 
     async def IsReady(

pynumaflow/sideinput/servicer/servicer.py

@@ -2,7 +2,7 @@
 
 from pynumaflow._constants import (
     _LOGGER,
-    ERR_SIDE_INPUT_RETRIEVAL_EXCEPTION,
+    ERR_UDF_EXCEPTION_STRING,
 )
 from pynumaflow.proto.sideinput import sideinput_pb2_grpc, sideinput_pb2
 from pynumaflow.shared.server import exit_on_error
@@ -28,7 +28,7 @@ def RetrieveSideInput(
         try:
             rspn = self.__retrieve_handler()
         except BaseException as err:
-            err_msg = f"{ERR_SIDE_INPUT_RETRIEVAL_EXCEPTION}: {repr(err)}"
+            err_msg = f"{ERR_UDF_EXCEPTION_STRING}: {repr(err)}"
             _LOGGER.critical(err_msg, exc_info=True)
             exit_on_error(context, err_msg)
             return

pynumaflow/sinker/servicer/async_servicer.py

@@ -13,7 +13,7 @@
     build_sink_resp_results,
 )
 from pynumaflow.types import NumaflowServicerContext
-from pynumaflow._constants import _LOGGER, STREAM_EOF, ERR_SINK_EXCEPTION
+from pynumaflow._constants import _LOGGER, STREAM_EOF, ERR_UDF_EXCEPTION_STRING
 
 
 class AsyncSinkServicer(sink_pb2_grpc.SinkServicer):
@@ -85,7 +85,7 @@ async def SinkFn(
             # if there is an exception, we will mark all the responses as a failure
             err_msg = f"UDSinkError: {repr(err)}"
             _LOGGER.critical(err_msg, exc_info=True)
-            await handle_async_error(context, err, ERR_SINK_EXCEPTION)
+            await handle_async_error(context, err, ERR_UDF_EXCEPTION_STRING)
             return
 
     async def __invoke_sink(

pynumaflow/sourcer/servicer/async_servicer.py

@@ -11,7 +11,7 @@
 from pynumaflow.proto.sourcer import source_pb2
 from pynumaflow.proto.sourcer import source_pb2_grpc
 from pynumaflow.types import NumaflowServicerContext
-from pynumaflow._constants import _LOGGER, STREAM_EOF, ERR_SOURCE_EXCEPTION
+from pynumaflow._constants import _LOGGER, STREAM_EOF, ERR_UDF_EXCEPTION_STRING
 
 
 def _create_read_handshake_response():
@@ -119,7 +119,7 @@ async def ReadFn(
                 yield _create_eot_response()
         except BaseException as err:
             _LOGGER.critical("User-Defined Source ReadFn error", exc_info=True)
-            await handle_async_error(context, err, ERR_SOURCE_EXCEPTION)
+            await handle_async_error(context, err, ERR_UDF_EXCEPTION_STRING)
 
     async def __invoke_read(self, req, niter):
         """Invoke the read handler and manage the iterator."""
@@ -165,7 +165,7 @@ async def AckFn(
                 yield _create_ack_response()
         except BaseException as err:
             _LOGGER.critical("User-Defined Source AckFn error", exc_info=True)
-            await handle_async_error(context, err, ERR_SOURCE_EXCEPTION)
+            await handle_async_error(context, err, ERR_UDF_EXCEPTION_STRING)
 
     async def IsReady(
         self, request: _empty_pb2.Empty, context: NumaflowServicerContext
@@ -187,7 +187,7 @@ async def PendingFn(
             count = await self.__source_pending_handler()
         except BaseException as err:
             _LOGGER.critical("PendingFn Error", exc_info=True)
-            await handle_async_error(context, err, ERR_SOURCE_EXCEPTION)
+            await handle_async_error(context, err, ERR_UDF_EXCEPTION_STRING)
             return
         resp = source_pb2.PendingResponse.Result(count=count.count)
         return source_pb2.PendingResponse(result=resp)
@@ -202,7 +202,7 @@ async def PartitionsFn(
             partitions = await self.__source_partitions_handler()
         except BaseException as err:
             _LOGGER.critical("PartitionsFn Error", exc_info=True)
-            await handle_async_error(context, err, ERR_SOURCE_EXCEPTION)
+            await handle_async_error(context, err, ERR_UDF_EXCEPTION_STRING)
             return
         resp = source_pb2.PartitionsResponse.Result(partitions=partitions.partitions)
         return source_pb2.PartitionsResponse(result=resp)

pynumaflow/sourcetransformer/servicer/_servicer.py

@@ -16,7 +16,7 @@
     _LOGGER,
     STREAM_EOF,
     NUM_THREADS_DEFAULT,
-    ERR_TRANSFORMER_EXCEPTION,
+    ERR_UDF_EXCEPTION_STRING,
 )
 
 
@@ -79,7 +79,7 @@ def SourceTransformFn(
                 if isinstance(res, BaseException):
                     # Terminate the current server process due to exception
                     exit_on_error(
-                        context, f"{ERR_TRANSFORMER_EXCEPTION}: {repr(res)}", parent=self.multiproc
+                        context, f"{ERR_UDF_EXCEPTION_STRING}: {repr(res)}", parent=self.multiproc
                     )
                     return
                 # return the result
@@ -93,7 +93,7 @@ def SourceTransformFn(
             _LOGGER.critical("UDFError, re-raising the error", exc_info=True)
             # Terminate the current server process due to exception
             exit_on_error(
-                context, f"{ERR_TRANSFORMER_EXCEPTION}: {repr(err)}", parent=self.multiproc
+                context, f"{ERR_UDF_EXCEPTION_STRING}: {repr(err)}", parent=self.multiproc
             )
             return
 

tests/sink/test_async_sink.py

@@ -11,6 +11,7 @@
 
 from pynumaflow import setup_logging
 from pynumaflow._constants import (
+    ERR_UDF_EXCEPTION_STRING,
     UD_CONTAINER_FALLBACK_SINK,
     FALLBACK_SINK_SOCK_PATH,
     FALLBACK_SINK_SERVER_INFO_FILE_PATH,
@@ -183,7 +184,7 @@ def test_sink_err(self) -> None:
                 pass
         except BaseException as e:
             self.assertTrue(
-                "UDF_EXECUTION_ERROR(sink): ValueError('test_mock_err_message')" in e.__str__()
+                f"{ERR_UDF_EXCEPTION_STRING}: ValueError('test_mock_err_message')" in e.__str__()
             )
             return
         except grpc.RpcError as e:

tests/source/test_async_source_err.py

@@ -9,6 +9,7 @@
 from grpc.aio._server import Server
 
 from pynumaflow import setup_logging
+from pynumaflow._constants import ERR_UDF_EXCEPTION_STRING
 from pynumaflow.sourcer import SourceAsyncServer
 from pynumaflow.proto.sourcer import source_pb2_grpc
 from google.protobuf import empty_pb2 as _empty_pb2
@@ -93,7 +94,7 @@ def test_read_error(self) -> None:
                     pass
             except BaseException as e:
                 self.assertTrue(
-                    "UDF_EXECUTION_ERROR(source): TypeError("
+                    f"{ERR_UDF_EXCEPTION_STRING}: TypeError("
                     '"handle_async_error() missing 1 required positional argument: '
                     "'exception_type'\")" in e.__str__()
                 )