Invocation errors and refactoring of Function Executor and Executor protocols (#1499)

* Implement InvocationError

If InvocationError is raised by customer function code then
then its current invocation fails. InvocationError requires
a string message argument that can be retrieved via Server API
and Tensrolake CLI.

Also refactor FunctionExecutor and Executor grpc protocols
because we don't need backward compatibility right now - it's
a good window to do this.

The CI tests will fail until we deploy the required changes
in Indexify (due to FE and Executor grpc protocol changes).

* Remove Server side routing fallback

Now Function Executor always provides next_functions in its run task
response. The next functions are either produced by the function
or edges from the graph definition object. This allows FE to fully
control next_functions and server just following with the decisions.
This moves the next_functions list ownership to a single place (FE/SDK)
and also simplifies the FE and Executor protocols.

* Use DataPayload for InvocationError message/payload

This makes it much more flexible wrt the size of the InvocationError
message and lets SDK own what can be stored there.
Users can now add any size of string message to InvocationError
and in the future we can easily allow DataModels there and etc
if needed.

This also removes the need to store the InvocationError message
in Server state store. All we need is just put DataPayload there.

* Implement InvocationError on Server side

Ingest invocation error message payload, fail graph invocation
if invocation error happened, add invocation error details to
HTTP API invocation object.

* Use pypi tensorlake instead of local

This was used for local development, shouldn't be merged.

* Add Executor proto enum comments and Server enum comments

Author: eabatalov

indexify/src/indexify/executor/executor.py

@@ -52,7 +52,7 @@ def __init__(
         blob_store: BLOBStore,
         host_resources_provider: HostResourcesProvider,
     ):
-        self._logger = structlog.get_logger(module=__name__)
+        self._logger = structlog.get_logger(module=__name__, executor_id=id)
         protocol: str = "http"
         if config_path:
             self._logger.info("running the extractor with TLS enabled")

indexify/src/indexify/executor/function_executor/function_executor.py

@@ -7,6 +7,8 @@
 from tensorlake.function_executor.proto.function_executor_pb2 import (
     InfoRequest,
     InfoResponse,
+    InitializationFailureReason,
+    InitializationOutcomeCode,
     InitializeRequest,
     InitializeResponse,
 )
@@ -315,16 +317,42 @@ async def _initialize_server(
                 initialize_request,
                 timeout=customer_code_timeout_sec,
             )
-            # TODO: set real stdout and stderr when their proper capturing on FE initialization is implemented.
-            if initialize_response.success:
-                return FunctionExecutorInitializationResult()
-            elif initialize_response.HasField("customer_error"):
+
+            if (
+                initialize_response.outcome_code
+                == InitializationOutcomeCode.INITIALIZE_OUTCOME_CODE_SUCCESS
+            ):
                 return FunctionExecutorInitializationResult(
-                    error=FunctionExecutorInitializationError.FUNCTION_ERROR,
-                    stderr=initialize_response.customer_error,
+                    stdout=initialize_response.stdout, stderr=initialize_response.stderr
                 )
+            elif (
+                initialize_response.outcome_code
+                == InitializationOutcomeCode.INITIALIZE_OUTCOME_CODE_FAILURE
+            ):
+                if (
+                    initialize_response.failure_reason
+                    == InitializationFailureReason.INITIALIZATION_FAILURE_REASON_FUNCTION_ERROR
+                ):
+                    return FunctionExecutorInitializationResult(
+                        error=FunctionExecutorInitializationError.FUNCTION_ERROR,
+                        stdout=initialize_response.stdout,
+                        stderr=initialize_response.stderr,
+                    )
+                elif (
+                    initialize_response.failure_reason
+                    == InitializationFailureReason.INITIALIZATION_FAILURE_REASON_INTERNAL_ERROR
+                ):
+                    # Don't add stdout/stderr because this is customer data.
+                    raise RuntimeError("initialize RPC failed with internal error")
+                else:
+                    raise ValueError(
+                        f"unexpected failure reason {InitializationFailureReason.Name(initialize_response.failure_reason)} in initialize RPC response"
+                    )
             else:
-                raise Exception("initialize RPC failed at function executor server")
+                raise ValueError(
+                    f"unexpected outcome code {InitializationOutcomeCode.Name(initialize_response.outcome_code)} in initialize RPC response"
+                )
+
         except grpc.aio.AioRpcError as e:
             if e.code() == grpc.StatusCode.DEADLINE_EXCEEDED:
                 return FunctionExecutorInitializationResult(

indexify/src/indexify/executor/function_executor/invocation_state_client.py

@@ -8,6 +8,7 @@
     InvocationStateRequest,
     InvocationStateResponse,
     SerializedObject,
+    SerializedObjectEncoding,
     SetInvocationStateResponse,
 )
 from tensorlake.function_executor.proto.function_executor_pb2_grpc import (
@@ -25,6 +26,10 @@
     metric_server_set_state_requests,
 )
 
+# We're currently only supporting CloudPickle for invocation state values.
+# FIXME: if Function Executor sends us something else then we fail the calls.
+_VALUE_CONTENT_TYPE = "application/octet-stream"
+
 
 class InvocationStateClient:
     """InvocationStateClient is a client for the invocation state server of a Function Executor.
@@ -196,14 +201,21 @@ async def _set_server_state(
         url: str = (
             f"{self._base_url}/internal/namespaces/{self._namespace}/compute_graphs/{self._graph}/invocations/{invocation_id}/ctx/{key}"
         )
-        payload = value.bytes if value.HasField("bytes") else value.string
+        if (
+            value.encoding
+            != SerializedObjectEncoding.SERIALIZED_OBJECT_ENCODING_BINARY_PICKLE
+        ):
+            raise ValueError(
+                f"Unsupported value encoding: {SerializedObjectEncoding.Name(value.encoding)}. "
+                "Only binary pickle is supported for invocation state values."
+            )
 
         response = await self._http_client.post(
             url=url,
             files=[
                 (
                     "value",
-                    ("value", payload, value.content_type),
+                    ("value", value.data, _VALUE_CONTENT_TYPE),
                 ),
             ],
         )
@@ -245,7 +257,7 @@ async def _get_server_state(
             )
             raise
 
-        return serialized_object_from_http_response(response)
+        return _serialized_object_from_http_response(response)
 
     def _validate_request(self, request: InvocationStateRequest) -> None:
         (
@@ -265,17 +277,15 @@ def _validate_request(self, request: InvocationStateRequest) -> None:
             raise ValueError("unknown request type")
 
 
-def serialized_object_from_http_response(response: httpx.Response) -> SerializedObject:
-    # We're hardcoding the content type currently used by Python SDK. It might change in the future.
-    # There's no other way for now to determine if the response is a bytes or string.
-    if response.headers["content-type"] in [
-        "application/octet-stream",
-        "application/pickle",
-    ]:
-        return SerializedObject(
-            bytes=response.content, content_type=response.headers["content-type"]
-        )
-    else:
-        return SerializedObject(
-            string=response.text, content_type=response.headers["content-type"]
+def _serialized_object_from_http_response(response: httpx.Response) -> SerializedObject:
+    if response.headers["content-type"] != _VALUE_CONTENT_TYPE:
+        raise ValueError(
+            f"Unexpected content type: {response.headers['content-type']}. "
+            f"Expected: {_VALUE_CONTENT_TYPE}."
         )
+
+    return SerializedObject(
+        data=response.content,
+        encoding=SerializedObjectEncoding.SERIALIZED_OBJECT_ENCODING_BINARY_PICKLE,
+        encoding_version=0,
+    )

indexify/src/indexify/executor/function_executor_controller/completed_task_metrics.py

@@ -66,6 +66,7 @@ def emit_completed_task_metrics(task_info: TaskInfo, logger: Any) -> None:
         elif task_failure_reason in [
             TaskFailureReason.TASK_FAILURE_REASON_FUNCTION_ERROR,
             TaskFailureReason.TASK_FAILURE_REASON_FUNCTION_TIMEOUT,
+            TaskFailureReason.TASK_FAILURE_REASON_INVOCATION_ERROR,
         ]:
             metric_tasks_completed.labels(
                 outcome_code=METRIC_TASKS_COMPLETED_OUTCOME_CODE_FAILURE,

indexify/src/indexify/executor/function_executor_controller/create_function_executor.py

@@ -5,6 +5,7 @@
 from tensorlake.function_executor.proto.function_executor_pb2 import (
     InitializeRequest,
     SerializedObject,
+    SerializedObjectEncoding,
 )
 
 from indexify.executor.blob_store.blob_store import BLOBStore

indexify/src/indexify/executor/function_executor_controller/downloads.py

@@ -4,7 +4,10 @@
 from typing import Any, Optional
 
 import nanoid
-from tensorlake.function_executor.proto.function_executor_pb2 import SerializedObject
+from tensorlake.function_executor.proto.function_executor_pb2 import (
+    SerializedObject,
+    SerializedObjectEncoding,
+)
 
 from indexify.executor.blob_store.blob_store import BLOBStore
 from indexify.proto.executor_api_pb2 import (
@@ -179,20 +182,28 @@ def _serialized_object_from_data_payload_proto(
     """
     if data_payload.encoding == DataPayloadEncoding.DATA_PAYLOAD_ENCODING_BINARY_PICKLE:
         return SerializedObject(
-            bytes=data,
-            content_type="application/octet-stream",
+            data=data,
+            encoding=SerializedObjectEncoding.SERIALIZED_OBJECT_ENCODING_BINARY_PICKLE,
+            encoding_version=data_payload.encoding_version,
         )
     elif data_payload.encoding == DataPayloadEncoding.DATA_PAYLOAD_ENCODING_UTF8_TEXT:
         return SerializedObject(
-            content_type="text/plain",
-            string=data.decode("utf-8"),
+            data=data,
+            encoding=SerializedObjectEncoding.SERIALIZED_OBJECT_ENCODING_UTF8_TEXT,
+            encoding_version=data_payload.encoding_version,
         )
     elif data_payload.encoding == DataPayloadEncoding.DATA_PAYLOAD_ENCODING_UTF8_JSON:
-        result = SerializedObject(
-            content_type="application/json",
-            string=data.decode("utf-8"),
+        return SerializedObject(
+            data=data,
+            encoding=SerializedObjectEncoding.SERIALIZED_OBJECT_ENCODING_UTF8_JSON,
+            encoding_version=data_payload.encoding_version,
+        )
+    elif data_payload.encoding == DataPayloadEncoding.DATA_PAYLOAD_ENCODING_BINARY_ZIP:
+        return SerializedObject(
+            data=data,
+            encoding=SerializedObjectEncoding.SERIALIZED_OBJECT_ENCODING_BINARY_ZIP,
+            encoding_version=data_payload.encoding_version,
         )
-        return result
 
     raise ValueError(
         f"Can't convert data payload {data_payload} into serialized object"

indexify/src/indexify/executor/function_executor_controller/function_executor_controller.py

@@ -197,7 +197,6 @@ def remove_task(self, task_id: str) -> None:
         task_info.is_cancelled = True
         logger.info(
             "cancelling task",
-            allocation_id=task_info.allocation.allocation_id,
         )
         if task_info.aio_task is not None:
             task_info.aio_task.cancel()
@@ -332,7 +331,7 @@ async def _control_loop(self) -> None:
                     self._logger.error(
                         "unexpected exception in function executor controller control loop",
                         exc_info=e,
-                        fe_event=str(event),
+                        event_type=event.event_type.name,
                     )
 
     def _handle_event(self, event: BaseEvent) -> None:
@@ -786,16 +785,14 @@ def _to_task_result_proto(output: TaskOutput) -> TaskResult:
         graph_invocation_id=output.allocation.task.graph_invocation_id,
         reducer=output.reducer,
         outcome_code=output.outcome_code,
-        next_functions=(output.router_output.edges if output.router_output else []),
+        failure_reason=output.failure_reason,
+        next_functions=output.next_functions,
         function_outputs=output.uploaded_data_payloads,
+        invocation_error_output=output.uploaded_invocation_error_output,
     )
-    if output.failure_reason is not None:
-        task_result.failure_reason = output.failure_reason
     if output.uploaded_stdout is not None:
         task_result.stdout.CopyFrom(output.uploaded_stdout)
     if output.uploaded_stderr is not None:
         task_result.stderr.CopyFrom(output.uploaded_stderr)
-    if output.router_output is not None:
-        task_result.routing.next_functions[:] = output.router_output.edges
 
     return task_result

indexify/src/indexify/executor/function_executor_controller/run_task.py

@@ -6,6 +6,13 @@
 from tensorlake.function_executor.proto.function_executor_pb2 import (
     RunTaskRequest,
     RunTaskResponse,
+    SerializedObject,
+)
+from tensorlake.function_executor.proto.function_executor_pb2 import (
+    TaskFailureReason as FETaskFailureReason,
+)
+from tensorlake.function_executor.proto.function_executor_pb2 import (
+    TaskOutcomeCode as FETaskOutcomeCode,
 )
 from tensorlake.function_executor.proto.function_executor_pb2_grpc import (
     FunctionExecutorStub,
@@ -83,6 +90,7 @@ async def run_task_on_function_executor(
         task_info.output = _task_output_from_function_executor_response(
             allocation=task_info.allocation,
             response=response,
+            logger=logger,
         )
     except grpc.aio.AioRpcError as e:
         if e.code() == grpc.StatusCode.DEADLINE_EXCEEDED:
@@ -140,45 +148,48 @@ async def run_task_on_function_executor(
 
 
 def _task_output_from_function_executor_response(
-    allocation: TaskAllocation, response: RunTaskResponse
+    allocation: TaskAllocation, response: RunTaskResponse, logger: Any
 ) -> TaskOutput:
     response_validator = MessageValidator(response)
     response_validator.required_field("stdout")
     response_validator.required_field("stderr")
     response_validator.required_field("is_reducer")
-    response_validator.required_field("success")
+    response_validator.required_field("outcome_code")
 
     metrics = TaskMetrics(counters={}, timers={})
     if response.HasField("metrics"):
         # Can be None if e.g. function failed.
         metrics.counters = dict(response.metrics.counters)
         metrics.timers = dict(response.metrics.timers)
 
-    output = TaskOutput(
+    outcome_code: TaskOutcomeCode = _to_task_outcome_code(
+        response.outcome_code, logger=logger
+    )
+    failure_reason: Optional[TaskFailureReason] = None
+    invocation_error_output: Optional[SerializedObject] = None
+
+    if outcome_code == TaskOutcomeCode.TASK_OUTCOME_CODE_FAILURE:
+        response_validator.required_field("failure_reason")
+        failure_reason: Optional[TaskFailureReason] = _to_task_failure_reason(
+            response.failure_reason, logger
+        )
+        if failure_reason == TaskFailureReason.TASK_FAILURE_REASON_INVOCATION_ERROR:
+            response_validator.required_field("invocation_error_output")
+            invocation_error_output = response.invocation_error_output
+
+    return TaskOutput(
         allocation=allocation,
-        outcome_code=(
-            TaskOutcomeCode.TASK_OUTCOME_CODE_SUCCESS
-            if response.success
-            else TaskOutcomeCode.TASK_OUTCOME_CODE_FAILURE
-        ),
-        failure_reason=(
-            None
-            if response.success
-            else TaskFailureReason.TASK_FAILURE_REASON_FUNCTION_ERROR
-        ),
+        outcome_code=outcome_code,
+        failure_reason=failure_reason,
+        invocation_error_output=invocation_error_output,
+        function_outputs=response.function_outputs,
+        next_functions=response.next_functions,
         stdout=response.stdout,
         stderr=response.stderr,
         reducer=response.is_reducer,
         metrics=metrics,
     )
 
-    if response.HasField("function_output"):
-        output.function_output = response.function_output
-    if response.HasField("router_output"):
-        output.router_output = response.router_output
-
-    return output
-
 
 def _log_task_execution_finished(output: TaskOutput, logger: Any) -> None:
     logger.info(
@@ -191,3 +202,40 @@ def _log_task_execution_finished(output: TaskOutput, logger: Any) -> None:
             else None
         ),
     )
+
+
+def _to_task_outcome_code(
+    fe_task_outcome_code: FETaskOutcomeCode, logger
+) -> TaskOutcomeCode:
+    if fe_task_outcome_code == FETaskOutcomeCode.TASK_OUTCOME_CODE_SUCCESS:
+        return TaskOutcomeCode.TASK_OUTCOME_CODE_SUCCESS
+    elif fe_task_outcome_code == FETaskOutcomeCode.TASK_OUTCOME_CODE_FAILURE:
+        return TaskOutcomeCode.TASK_OUTCOME_CODE_FAILURE
+    else:
+        logger.warning(
+            "Unknown TaskOutcomeCode received from Function Executor",
+            value=FETaskOutcomeCode.Name(fe_task_outcome_code),
+        )
+        return TaskOutcomeCode.TASK_OUTCOME_CODE_UNKNOWN
+
+
+def _to_task_failure_reason(
+    fe_task_failure_reason: FETaskFailureReason, logger: Any
+) -> TaskFailureReason:
+    if fe_task_failure_reason == FETaskFailureReason.TASK_FAILURE_REASON_FUNCTION_ERROR:
+        return TaskFailureReason.TASK_FAILURE_REASON_FUNCTION_ERROR
+    elif (
+        fe_task_failure_reason
+        == FETaskFailureReason.TASK_FAILURE_REASON_INVOCATION_ERROR
+    ):
+        return TaskFailureReason.TASK_FAILURE_REASON_INVOCATION_ERROR
+    elif (
+        fe_task_failure_reason == FETaskFailureReason.TASK_FAILURE_REASON_INTERNAL_ERROR
+    ):
+        return TaskFailureReason.TASK_FAILURE_REASON_INTERNAL_ERROR
+    else:
+        logger.warning(
+            "Unknown TaskFailureReason received from Function Executor",
+            value=FETaskFailureReason.Name(fe_task_failure_reason),
+        )
+        return TaskFailureReason.TASK_FAILURE_REASON_UNKNOWN