indexify 0.2.40__py3-none-any.whl → 0.2.41__py3-none-any.whl

indexify/executor/task_reporter.py

@@ -9,7 +9,6 @@ from indexify.common_util import get_httpx_client
 from indexify.executor.api_objects import RouterOutput as ApiRouterOutput
 from indexify.executor.api_objects import TaskResult
 from indexify.executor.task_store import CompletedTask
-from indexify.functions_sdk.object_serializer import get_serializer
 
 logger = structlog.get_logger(__name__)
 
@@ -42,19 +41,20 @@ class TaskReporter:
         self._client = get_httpx_client(config_path)
 
     def report_task_outcome(self, completed_task: CompletedTask):
-
         report = ReportingData()
         fn_outputs = []
-        for output in completed_task.outputs or []:
-            serializer = get_serializer(output.encoder)
-            fn_outputs.append(
-                (
-                    "node_outputs",
-                    (nanoid.generate(), output.payload, serializer.content_type),
+
+        if completed_task.function_output:
+            for output in completed_task.function_output.outputs or []:
+                payload = output.bytes if output.HasField("bytes") else output.string
+                fn_outputs.append(
+                    (
+                        "node_outputs",
+                        (nanoid.generate(), payload, output.content_type),
+                    )
                 )
-            )
-            report.output_count += 1
-            report.output_total_bytes += len(output.payload)
+                report.output_count += 1
+                report.output_total_bytes += len(payload)
 
         if completed_task.stdout:
             fn_outputs.append(
@@ -134,24 +134,17 @@ class TaskReporter:
             kwargs["files"] = fn_outputs
         else:
             kwargs["files"] = FORCE_MULTIPART
-        try:
-            response = self._client.post(
-                url=f"{self._base_url}/internal/ingest_files",
-                **kwargs,
-            )
-        except Exception as e:
-            logger.error(
-                "failed to report task outcome",
-                task_id=completed_task.task.id,
-                retries=completed_task.reporting_retries,
-                error=type(e).__name__,
-                message=str(e),
-            )
-            raise e
+
+        response = self._client.post(
+            url=f"{self._base_url}/internal/ingest_files",
+            **kwargs,
+        )
 
         try:
             response.raise_for_status()
         except Exception as e:
+            # Caller catches and logs the exception.
+            # Log response details here for easier debugging.
             logger.error(
                 "failed to report task outcome",
                 task_id=completed_task.task.id,

indexify/executor/task_store.py

@@ -1,23 +1,38 @@
 import asyncio
 from typing import Dict, List, Literal, Optional
 
-from pydantic import BaseModel
-from rich import print
+import structlog
 
-from indexify.functions_sdk.data_objects import IndexifyData, RouterOutput
+from indexify.function_executor.proto.function_executor_pb2 import (
+    FunctionOutput,
+    RouterOutput,
+)
 
 from .api_objects import Task
 
-
-class CompletedTask(BaseModel):
-    task: Task
-    task_outcome: Literal["success", "failure"]
-    outputs: Optional[List[IndexifyData]] = None
-    router_output: Optional[RouterOutput] = None
-    stdout: Optional[str] = None
-    stderr: Optional[str] = None
-    reducer: bool = False
-    reporting_retries: int = 0
+logger = structlog.get_logger(module=__name__)
+
+
+class CompletedTask:
+    def __init__(
+        self,
+        task: Task,
+        task_outcome: Literal["success", "failure"],
+        function_output: Optional[FunctionOutput] = None,
+        router_output: Optional[RouterOutput] = None,
+        stdout: Optional[str] = None,
+        stderr: Optional[str] = None,
+        reducer: bool = False,
+        reporting_retries: int = 0,
+    ):
+        self.task = task
+        self.task_outcome = task_outcome
+        self.function_output = function_output
+        self.router_output = router_output
+        self.stdout = stdout
+        self.stderr = stderr
+        self.reducer = reducer
+        self.reporting_retries = reporting_retries
 
 
 class TaskStore:
@@ -41,8 +56,12 @@ class TaskStore:
                 or (task.id in self._finished)
             ):
                 continue
-            print(
-                f"[bold] task store: [/bold] added task: {task.id} graph: {task.compute_graph} fn: {task.compute_fn} to queue"
+            logger.info(
+                "added task",
+                task_id=task.id,
+                namespace=task.namespace,
+                graph=task.compute_graph,
+                fn=task.compute_fn,
             )
             self._tasks[task.id] = task
             self._new_task_event.set()
@@ -87,7 +106,7 @@ class TaskStore:
     def mark_reported(self, task_id: str):
         self._tasks.pop(task_id)
         self._finished.pop(task_id)
-        print(f"[bold] task store: [/bold] removed task: {task_id} from queue")
+        logger.info("removed task", task_id=task_id)
 
     def report_failed(self, task_id: str):
         if self._finished[task_id].task_outcome != "Failed":

indexify/function_executor/function_executor_service.py

@@ -0,0 +1,86 @@
+from typing import Optional, Union
+
+import grpc
+import structlog
+
+from indexify.functions_sdk.indexify_functions import (
+    IndexifyFunction,
+    IndexifyRouter,
+)
+from indexify.functions_sdk.object_serializer import get_serializer
+
+from .handlers.run_function.handler import Handler as RunTaskHandler
+from .handlers.run_function.request_validator import (
+    RequestValidator as RunTaskRequestValidator,
+)
+from .initialize_request_validator import InitializeRequestValidator
+from .proto.function_executor_pb2 import (
+    InitializeRequest,
+    InitializeResponse,
+    RunTaskRequest,
+    RunTaskResponse,
+)
+from .proto.function_executor_pb2_grpc import FunctionExecutorServicer
+
+
+class FunctionExecutorService(FunctionExecutorServicer):
+    def __init__(self, indexify_server_address: str, config_path: Optional[str]):
+        self._indexify_server_address = indexify_server_address
+        self._config_path = config_path
+        self._logger = structlog.get_logger(module=__name__)
+        self._namespace: Optional[str] = None
+        self._graph_name: Optional[str] = None
+        self._graph_version: Optional[int] = None
+        self._function_name: Optional[str] = None
+        self._function: Optional[Union[IndexifyFunction, IndexifyRouter]] = None
+
+    def initialize(
+        self, request: InitializeRequest, context: grpc.ServicerContext
+    ) -> InitializeResponse:
+        request_validator: InitializeRequestValidator = InitializeRequestValidator(
+            request
+        )
+        request_validator.check()
+
+        self._namespace = request.namespace
+        self._graph_name = request.graph_name
+        self._graph_version = request.graph_version
+        self._function_name = request.function_name
+        # The function is only loaded once per Function Executor. It's important to use a single
+        # loaded function so all the tasks when executed are sharing the same memory. This allows
+        # implementing smart caching in customer code. E.g. load a model into GPU only once and
+        # share the model's file descriptor between all tasks or download function configuration
+        # only once.
+        graph_serializer = get_serializer(request.graph.content_type)
+        graph = graph_serializer.deserialize(request.graph.bytes)
+        self._function = graph_serializer.deserialize(graph[request.function_name])
+
+        self._logger = self._logger.bind(
+            namespace=request.namespace,
+            graph_name=request.graph_name,
+            graph_version=str(request.graph_version),
+            function_name=request.function_name,
+        )
+        self._logger.info("initialized function executor service")
+
+        return InitializeResponse(success=True)
+
+    def run_task(
+        self, request: RunTaskRequest, context: grpc.ServicerContext
+    ) -> RunTaskResponse:
+        # Customer function code never raises an exception because we catch all of them and add
+        # their details to the response. We can only get an exception here if our own code failed.
+        # If our code raises an exception the grpc framework converts it into GRPC_STATUS_UNKNOWN
+        # error with the exception message. Differentiating errors is not needed for now.
+        RunTaskRequestValidator(request=request).check()
+        return RunTaskHandler(
+            request=request,
+            namespace=self._namespace,
+            graph_name=self._graph_name,
+            graph_version=self._graph_version,
+            function_name=self._function_name,
+            function=self._function,
+            logger=self._logger,
+            indexify_server_addr=self._indexify_server_address,
+            config_path=self._config_path,
+        ).run()

indexify/function_executor/handlers/run_function/function_inputs_loader.py

@@ -0,0 +1,54 @@
+from typing import Optional
+
+from pydantic import BaseModel
+
+from indexify.function_executor.proto.function_executor_pb2 import (
+    RunTaskRequest,
+    SerializedObject,
+)
+from indexify.functions_sdk.data_objects import IndexifyData
+from indexify.functions_sdk.object_serializer import get_serializer
+
+
+class FunctionInputs(BaseModel):
+    input: IndexifyData
+    init_value: Optional[IndexifyData] = None
+
+
+class FunctionInputsLoader:
+    def __init__(self, request: RunTaskRequest):
+        self._request = request
+
+    def load(self) -> FunctionInputs:
+        return FunctionInputs(
+            input=self._function_input(),
+            init_value=self._accumulator_input(),
+        )
+
+    def _function_input(self) -> IndexifyData:
+        return _to_indexify_data(
+            self._request.graph_invocation_id, self._request.function_input
+        )
+
+    def _accumulator_input(self) -> Optional[IndexifyData]:
+        return (
+            _to_indexify_data(
+                self._request.graph_invocation_id, self._request.function_init_value
+            )
+            if self._request.HasField("function_init_value")
+            else None
+        )
+
+
+def _to_indexify_data(
+    input_id: str, serialized_object: SerializedObject
+) -> IndexifyData:
+    return IndexifyData(
+        input_id=input_id,
+        payload=(
+            serialized_object.bytes
+            if serialized_object.HasField("bytes")
+            else serialized_object.string
+        ),
+        encoder=get_serializer(serialized_object.content_type).encoding_type,
+    )

indexify/function_executor/handlers/run_function/handler.py

@@ -0,0 +1,149 @@
+import io
+import sys
+import traceback
+from contextlib import redirect_stderr, redirect_stdout
+from typing import Any, Optional, Union
+
+from indexify.function_executor.proto.function_executor_pb2 import (
+    RunTaskRequest,
+    RunTaskResponse,
+)
+from indexify.functions_sdk.indexify_functions import (
+    FunctionCallResult,
+    GraphInvocationContext,
+    IndexifyFunction,
+    IndexifyFunctionWrapper,
+    IndexifyRouter,
+    RouterCallResult,
+)
+from indexify.http_client import IndexifyClient
+
+from .function_inputs_loader import FunctionInputs, FunctionInputsLoader
+from .response_helper import ResponseHelper
+
+
+class Handler:
+    def __init__(
+        self,
+        request: RunTaskRequest,
+        namespace: str,
+        graph_name: str,
+        graph_version: int,
+        function_name: str,
+        function: Union[IndexifyFunction, IndexifyRouter],
+        logger: Any,
+        indexify_server_addr: str,
+        config_path: Optional[str],
+    ):
+        self._function_name: str = function_name
+        self._logger = logger.bind(
+            graph_invocation_id=request.graph_invocation_id,
+            task_id=request.task_id,
+        )
+        self._input_loader = FunctionInputsLoader(request)
+        self._response_helper = ResponseHelper(task_id=request.task_id)
+        # TODO: use files for stdout, stderr capturing. This puts a natural and thus reasonable
+        # rate limit on the rate of writes and allows to not consume expensive memory for function logs.
+        self._func_stdout: io.StringIO = io.StringIO()
+        self._func_stderr: io.StringIO = io.StringIO()
+
+        self._function_wrapper: IndexifyFunctionWrapper = IndexifyFunctionWrapper(
+            indexify_function=function,
+            context=GraphInvocationContext(
+                invocation_id=request.graph_invocation_id,
+                graph_name=graph_name,
+                graph_version=str(graph_version),
+                indexify_client=_indexify_client(
+                    logger=self._logger,
+                    namespace=namespace,
+                    indexify_server_addr=indexify_server_addr,
+                    config_path=config_path,
+                ),
+            ),
+        )
+
+    def run(self) -> RunTaskResponse:
+        """Runs the task.
+
+        Raises an exception if our own code failed, customer function failure doesn't result in any exception.
+        Details of customer function failure are returned in the response.
+        """
+        self._logger.info("running function")
+        inputs: FunctionInputs = self._input_loader.load()
+        self._flush_logs()
+        return self._run_func_safe_and_captured(inputs)
+
+    def _run_func_safe_and_captured(self, inputs: FunctionInputs) -> RunTaskResponse:
+        """Runs the customer function while capturing what happened in it.
+
+        Function stdout and stderr are captured so they don't get into Function Executor process stdout
+        and stderr. Never throws an Exception. Caller can determine if the function succeeded
+        using the response.
+        """
+        try:
+            with redirect_stdout(self._func_stdout), redirect_stderr(self._func_stderr):
+                return self._run_func(inputs)
+        except Exception:
+            return self._response_helper.failure_response(
+                message=traceback.format_exc(),
+                stdout=self._func_stdout.getvalue(),
+                stderr=self._func_stderr.getvalue(),
+            )
+
+    def _run_func(self, inputs: FunctionInputs) -> RunTaskResponse:
+        if _is_router(self._function_wrapper):
+            result: RouterCallResult = self._function_wrapper.invoke_router(
+                self._function_name, inputs.input
+            )
+            return self._response_helper.router_response(
+                result=result,
+                stdout=self._func_stdout.getvalue(),
+                stderr=self._func_stderr.getvalue(),
+            )
+        else:
+            result: FunctionCallResult = self._function_wrapper.invoke_fn_ser(
+                self._function_name, inputs.input, inputs.init_value
+            )
+            return self._response_helper.function_response(
+                result=result,
+                is_reducer=_func_is_reducer(self._function_wrapper),
+                stdout=self._func_stdout.getvalue(),
+                stderr=self._func_stderr.getvalue(),
+            )
+
+    def _flush_logs(self) -> None:
+        # Flush any logs buffered in memory before running the function with stdout, stderr capture.
+        # Otherwise our logs logged before this point will end up in the function's stdout.
+        # structlog.PrintLogger uses print function. This is why flushing with print works.
+        print("", flush=True)
+        sys.stdout.flush()
+        sys.stderr.flush()
+
+
+def _indexify_client(
+    logger: Any,
+    namespace: str,
+    indexify_server_addr: str,
+    config_path: Optional[str],
+) -> IndexifyClient:
+    # This client is required to implement key/value store functionality for customer functions.
+    protocol: str = "http"
+    if config_path:
+        logger.info("TLS is enabled")
+        protocol = "https"
+    return IndexifyClient(
+        service_url=f"{protocol}://{indexify_server_addr}",
+        namespace=namespace,
+        config_path=config_path,
+    )
+
+
+def _is_router(func_wrapper: IndexifyFunctionWrapper) -> bool:
+    return (
+        str(type(func_wrapper.indexify_function))
+        == "<class 'indexify.functions_sdk.indexify_functions.IndexifyRouter'>"
+    )
+
+
+def _func_is_reducer(func_wrapper: IndexifyFunctionWrapper) -> bool:
+    return func_wrapper.indexify_function.accumulate is not None

indexify/function_executor/handlers/run_function/request_validator.py

@@ -0,0 +1,24 @@
+from typing import Any
+
+from indexify.function_executor.proto.function_executor_pb2 import (
+    RunTaskRequest,
+)
+from indexify.function_executor.proto.message_validator import MessageValidator
+
+
+class RequestValidator:
+    def __init__(self, request: RunTaskRequest):
+        self._request = request
+        self._message_validator = MessageValidator(request)
+
+    def check(self):
+        """Validates the request.
+
+        Raises: ValueError: If the request is invalid.
+        """
+        (
+            self._message_validator.required_field("graph_invocation_id")
+            .required_field("task_id")
+            .required_serialized_object("function_input")
+            .optional_serialized_object("function_init_value")
+        )

indexify/function_executor/handlers/run_function/response_helper.py

@@ -0,0 +1,98 @@
+from typing import List
+
+from indexify.function_executor.proto.function_executor_pb2 import (
+    FunctionOutput,
+    RouterOutput,
+    RunTaskResponse,
+    SerializedObject,
+)
+from indexify.functions_sdk.data_objects import IndexifyData
+from indexify.functions_sdk.indexify_functions import (
+    FunctionCallResult,
+    RouterCallResult,
+)
+from indexify.functions_sdk.object_serializer import get_serializer
+
+
+class ResponseHelper:
+    """Helper class for generating RunFunctionResponse."""
+
+    def __init__(self, task_id: str):
+        self._task_id = task_id
+
+    def function_response(
+        self,
+        result: FunctionCallResult,
+        is_reducer: bool,
+        stdout: str = "",
+        stderr: str = "",
+    ) -> RunTaskResponse:
+        if result.traceback_msg is None:
+            return RunTaskResponse(
+                task_id=self._task_id,
+                function_output=self._to_function_output(result.ser_outputs),
+                router_output=None,
+                stdout=stdout,
+                stderr=stderr,
+                is_reducer=is_reducer,
+                success=True,
+            )
+        else:
+            return self.failure_response(
+                message=result.traceback_msg,
+                stdout=stdout,
+                stderr=stderr,
+            )
+
+    def router_response(
+        self,
+        result: RouterCallResult,
+        stdout: str = "",
+        stderr: str = "",
+    ) -> RunTaskResponse:
+        if result.traceback_msg is None:
+            return RunTaskResponse(
+                task_id=self._task_id,
+                function_output=None,
+                router_output=RouterOutput(edges=result.edges),
+                stdout=stdout,
+                stderr=stderr,
+                is_reducer=False,
+                success=True,
+            )
+        else:
+            return self.failure_response(
+                message=result.traceback_msg,
+                stdout=stdout,
+                stderr=stderr,
+            )
+
+    def failure_response(
+        self, message: str, stdout: str, stderr: str
+    ) -> RunTaskResponse:
+        stderr = "\n".join([stderr, message])
+        return RunTaskResponse(
+            task_id=self._task_id,
+            function_output=None,
+            router_output=None,
+            stdout=stdout,
+            stderr=stderr,
+            is_reducer=False,
+            success=False,
+        )
+
+    def _to_function_output(self, outputs: List[IndexifyData]) -> FunctionOutput:
+        output = FunctionOutput(outputs=[])
+        for ix_data in outputs:
+            serialized_object: SerializedObject = SerializedObject(
+                content_type=get_serializer(ix_data.encoder).content_type,
+            )
+            if isinstance(ix_data.payload, bytes):
+                serialized_object.bytes = ix_data.payload
+            elif isinstance(ix_data.payload, str):
+                serialized_object.string = ix_data.payload
+            else:
+                raise ValueError(f"Unsupported payload type: {type(ix_data.payload)}")
+
+            output.outputs.append(serialized_object)
+        return output

indexify/function_executor/initialize_request_validator.py

@@ -0,0 +1,22 @@
+from indexify.function_executor.proto.message_validator import MessageValidator
+
+from .proto.function_executor_pb2 import InitializeRequest
+
+
+class InitializeRequestValidator:
+    def __init__(self, request: InitializeRequest):
+        self._request = request
+        self._message_validator = MessageValidator(request)
+
+    def check(self):
+        """Validates the request.
+
+        Raises: ValueError: If the request is invalid.
+        """
+        (
+            self._message_validator.required_field("namespace")
+            .required_field("graph_name")
+            .required_field("graph_version")
+            .required_field("function_name")
+            .required_serialized_object("graph")
+        )

indexify/function_executor/proto/configuration.py

@@ -0,0 +1,13 @@
+# We send function inputs and outputs over gRPC.
+# -1 means unlimited. We don't want to limit the size of data customers are using.
+# The effective max message size in this case is about 1.9 GB, see the max payload test.
+# This is due to internal hard gRPC limits. When we want to increase the message sizes
+# we'll have to implement chunking for large messages.
+_MAX_GRPC_MESSAGE_LENGTH = -1
+
+GRPC_SERVER_OPTIONS = [
+    ("grpc.max_receive_message_length", _MAX_GRPC_MESSAGE_LENGTH),
+    ("grpc.max_send_message_length", _MAX_GRPC_MESSAGE_LENGTH),
+]
+
+GRPC_CHANNEL_OPTIONS = GRPC_SERVER_OPTIONS

indexify/function_executor/proto/function_executor.proto

@@ -0,0 +1,70 @@
+
+
+syntax = "proto3";
+
+package function_executor_service;
+
+// The messages should not use any Python SDK objects. Only Function Executor implemented
+// in Python is allowed to import Python SDK to run customer functions. This ensures that
+// all the other components can be written in any language.
+
+message SerializedObject {
+    oneof data {
+        // Set bytes_data if the object is serialized as bytes.
+        bytes bytes = 1;
+        // Set string_data if the object is serialized as string.
+        string string = 2;
+    }
+    // The content type determines the serializer used to serialize the object.
+    optional string content_type = 3;
+}
+
+// InitializeRequest contains information about the function
+// that Function Executor is going to run the tasks for.
+message InitializeRequest {
+    optional string namespace = 1;
+    optional string graph_name = 2;
+    optional int32 graph_version = 3;
+    optional string function_name = 5;
+    optional SerializedObject graph = 7;
+}
+
+message InitializeResponse {
+    optional bool success = 1;
+}
+
+message FunctionOutput {
+    repeated SerializedObject outputs = 1;
+}
+
+message RouterOutput {
+    repeated string edges = 1;
+}
+
+message RunTaskRequest {
+    optional string graph_invocation_id = 4;
+    optional string task_id = 6;
+    optional SerializedObject function_input = 9;
+    optional SerializedObject function_init_value = 10;
+}
+
+message RunTaskResponse {
+    optional string task_id = 1;
+    optional FunctionOutput function_output = 2;
+    optional RouterOutput router_output = 3;
+    optional string stdout = 4;
+    optional string stderr = 5;
+    optional bool is_reducer = 6;
+    optional bool success = 7;
+}
+
+service FunctionExecutor {
+    // Initializes the Function Executor to run tasks
+    // for a particular function. This method is called only
+    // once per Function Executor as it can only run a single function.
+    // It should be called before calling RunTask for the function.
+    rpc initialize(InitializeRequest) returns (InitializeResponse);
+    // Executes the task defined in the request.
+    // Multiple tasks can be running in parallel.
+    rpc run_task(RunTaskRequest) returns (RunTaskResponse);
+}