indexify 0.2.47__py3-none-any.whl → 0.3.0__py3-none-any.whl

indexify/executor/function_executor/single_task_runner.py

@@ -12,7 +12,7 @@ from indexify.function_executor.proto.function_executor_pb2_grpc import (
 )
 
 from ..api_objects import Task
-from .function_executor import FunctionExecutor
+from .function_executor import CustomerError, FunctionExecutor
 from .function_executor_state import FunctionExecutorState
 from .server.function_executor_server_factory import (
     FunctionExecutorServerConfiguration,
@@ -49,8 +49,18 @@ class SingleTaskRunner:
         Raises an exception if an error occured."""
         self._state.check_locked()
 
+        if self._state.is_shutdown:
+            raise RuntimeError("Function Executor state is shutting down.")
+
         if self._state.function_executor is None:
-            self._state.function_executor = await self._create_function_executor()
+            try:
+                await self._create_function_executor()
+            except CustomerError as e:
+                return TaskOutput(
+                    task=self._task_input.task,
+                    stderr=str(e),
+                    success=False,
+                )
 
         return await self._run()
 
@@ -58,36 +68,37 @@ class SingleTaskRunner:
         function_executor: FunctionExecutor = FunctionExecutor(
             server_factory=self._factory, logger=self._logger
         )
-        try:
-            config: FunctionExecutorServerConfiguration = (
-                FunctionExecutorServerConfiguration(
-                    image_uri=self._task_input.task.image_uri,
-                )
-            )
-            initialize_request: InitializeRequest = InitializeRequest(
-                namespace=self._task_input.task.namespace,
-                graph_name=self._task_input.task.compute_graph,
-                graph_version=self._task_input.task.graph_version,
-                function_name=self._task_input.task.compute_fn,
-                graph=self._task_input.graph,
+        config: FunctionExecutorServerConfiguration = (
+            FunctionExecutorServerConfiguration(
+                image_uri=self._task_input.task.image_uri,
             )
+        )
+        initialize_request: InitializeRequest = InitializeRequest(
+            namespace=self._task_input.task.namespace,
+            graph_name=self._task_input.task.compute_graph,
+            graph_version=self._task_input.task.graph_version,
+            function_name=self._task_input.task.compute_fn,
+            graph=self._task_input.graph,
+        )
+
+        try:
             await function_executor.initialize(
                 config=config,
                 initialize_request=initialize_request,
                 base_url=self._base_url,
                 config_path=self._config_path,
             )
-            return function_executor
-        except Exception as e:
-            self._logger.error(
-                "failed to initialize function executor",
-                exc_info=e,
-            )
+            self._state.function_executor = function_executor
+        except Exception:
             await function_executor.destroy()
             raise
 
     async def _run(self) -> TaskOutput:
         request: RunTaskRequest = RunTaskRequest(
+            namespace=self._task_input.task.namespace,
+            graph_name=self._task_input.task.compute_graph,
+            graph_version=self._task_input.task.graph_version,
+            function_name=self._task_input.task.compute_fn,
             graph_invocation_id=self._task_input.task.invocation_id,
             task_id=self._task_input.task.id,
             function_input=self._task_input.input,
@@ -97,7 +108,9 @@ class SingleTaskRunner:
         channel: grpc.aio.Channel = self._state.function_executor.channel()
 
         async with _RunningTaskContextManager(
-            task_input=self._task_input, function_executor_state=self._state
+            invocation_id=self._task_input.task.invocation_id,
+            task_id=self._task_input.task.id,
+            function_executor_state=self._state,
         ):
             response: RunTaskResponse = await FunctionExecutorStub(channel).run_task(
                 request
@@ -109,16 +122,20 @@ class _RunningTaskContextManager:
     """Performs all the actions required before and after running a task."""
 
     def __init__(
-        self, task_input: TaskInput, function_executor_state: FunctionExecutorState
+        self,
+        invocation_id: str,
+        task_id: str,
+        function_executor_state: FunctionExecutorState,
     ):
-        self._task_input: TaskInput = task_input
+        self._invocation_id: str = invocation_id
+        self._task_id: str = task_id
         self._state: FunctionExecutorState = function_executor_state
 
     async def __aenter__(self):
         self._state.increment_running_tasks()
         self._state.function_executor.invocation_state_client().add_task_to_invocation_id_entry(
-            task_id=self._task_input.task.id,
-            invocation_id=self._task_input.task.invocation_id,
+            task_id=self._task_id,
+            invocation_id=self._invocation_id,
         )
         # Unlock the state so other tasks can act depending on it.
         self._state.lock.release()
@@ -128,7 +145,7 @@ class _RunningTaskContextManager:
         await self._state.lock.acquire()
         self._state.decrement_running_tasks()
         self._state.function_executor.invocation_state_client().remove_task_to_invocation_id_entry(
-            task_id=self._task_input.task.id
+            task_id=self._task_id
         )
 
 

indexify/executor/function_executor/task_input.py

@@ -1,8 +1,6 @@
 from typing import Optional
 
-from indexify.function_executor.proto.function_executor_pb2 import (
-    SerializedObject,
-)
+from indexify.function_executor.proto.function_executor_pb2 import SerializedObject
 
 from ..api_objects import Task
 

indexify/executor/task_fetcher.py

@@ -1,13 +1,13 @@
 import json
 from importlib.metadata import version
-from typing import AsyncGenerator, Optional
+from typing import AsyncGenerator, List, Optional
 
 import structlog
 from httpx_sse import aconnect_sse
 
-from indexify.common_util import get_httpx_client
+from indexify.utils.http_client import get_httpx_client
 
-from .api_objects import ExecutorMetadata, Task
+from .api_objects import ExecutorMetadata, FunctionURI, Task
 from .runtime_probes import ProbeInfo, RuntimeProbes
 
 
@@ -19,8 +19,7 @@ class TaskFetcher:
         protocol: str,
         indexify_server_addr: str,
         executor_id: str,
-        name_alias: Optional[str] = None,
-        image_hash: Optional[int] = None,
+        function_allowlist: Optional[List[FunctionURI]],
         config_path: Optional[str] = None,
     ):
         self._protocol: str = protocol
@@ -33,8 +32,7 @@ class TaskFetcher:
             id=executor_id,
             executor_version=version("indexify"),
             addr="",
-            image_name=probe_info.image_name if name_alias is None else name_alias,
-            image_hash=(probe_info.image_hash if image_hash is None else image_hash),
+            function_allowlist=function_allowlist,
             labels=probe_info.labels,
         )
 

indexify/executor/task_reporter.py

@@ -4,12 +4,10 @@ from typing import Any, List, Optional, Tuple
 import nanoid
 from httpx import Timeout
 
-from indexify.common_util import get_httpx_client
-from indexify.executor.api_objects import RouterOutput, Task, TaskResult
-from indexify.function_executor.proto.function_executor_pb2 import (
-    FunctionOutput,
-)
+from indexify.function_executor.proto.function_executor_pb2 import FunctionOutput
+from indexify.utils.http_client import get_httpx_client
 
+from .api_objects import RouterOutput, Task, TaskResult
 from .task_runner import TaskOutput
 
 

indexify/executor/task_runner.py

@@ -1,3 +1,4 @@
+import asyncio
 from typing import Any, Dict, Optional
 
 from .api_objects import Task
@@ -24,7 +25,9 @@ class TaskRunner:
         self._factory: FunctionExecutorServerFactory = function_executor_server_factory
         self._base_url: str = base_url
         self._config_path: Optional[str] = config_path
-        # We don't lock this map cause we never await while reading and modifying it.
+        # The fields below are protected by the lock.
+        self._lock: asyncio.Lock = asyncio.Lock()
+        self._is_shutdown: bool = False
         self._function_executor_states: Dict[str, FunctionExecutorState] = {}
 
     async def run(self, task_input: TaskInput, logger: Any) -> TaskOutput:
@@ -33,17 +36,31 @@ class TaskRunner:
             return await self._run(task_input, logger)
         except Exception as e:
             logger.error(
-                "failed running the task",
+                "failed running the task:",
                 exc_info=e,
             )
             return TaskOutput.internal_error(task_input.task)
 
     async def _run(self, task_input: TaskInput, logger: Any) -> TaskOutput:
-        state = self._get_or_create_state(task_input.task)
+        state = await self._get_or_create_state(task_input.task)
         async with state.lock:
             await self._run_task_policy(state, task_input.task)
             return await self._run_task(state, task_input, logger)
 
+    async def _get_or_create_state(self, task: Task) -> FunctionExecutorState:
+        async with self._lock:
+            if self._is_shutdown:
+                raise RuntimeError("Task runner is shutting down.")
+
+            id = _function_id_without_version(task)
+            if id not in self._function_executor_states:
+                state = FunctionExecutorState(
+                    function_id_with_version=_function_id_with_version(task),
+                    function_id_without_version=id,
+                )
+                self._function_executor_states[id] = state
+            return self._function_executor_states[id]
+
     async def _run_task_policy(self, state: FunctionExecutorState, task: Task) -> None:
         # Current policy for running tasks:
         #   - There can only be a single Function Executor per function regardless of function versions.
@@ -60,16 +77,6 @@ class TaskRunner:
             # At this point the state belongs to the version of the function from the task
             # and there are no running tasks in the Function Executor.
 
-    def _get_or_create_state(self, task: Task) -> FunctionExecutorState:
-        id = _function_id_without_version(task)
-        if id not in self._function_executor_states:
-            state = FunctionExecutorState(
-                function_id_with_version=_function_id_with_version(task),
-                function_id_without_version=id,
-            )
-            self._function_executor_states[id] = state
-        return self._function_executor_states[id]
-
     async def _run_task(
         self, state: FunctionExecutorState, task_input: TaskInput, logger: Any
     ) -> TaskOutput:
@@ -84,16 +91,16 @@ class TaskRunner:
         return await runner.run()
 
     async def shutdown(self) -> None:
-        # When shutting down there's no need to wait for completion of the running
-        # FunctionExecutor tasks.
-        while self._function_executor_states:
-            id, state = self._function_executor_states.popitem()
-            # At this point the state is not visible to new tasks.
-            # Only ongoing tasks who read it already have a reference to it.
-            await state.destroy_function_executor_not_locked()
-            # The task running inside the Function Executor will fail because it's destroyed.
-            # asyncio tasks waiting to run inside the Function Executor will get cancelled by
-            # the caller's shutdown code.
+        # Function Executors are outside the Executor process
+        # so they need to get cleaned up explicitly and reliably.
+        async with self._lock:
+            self._is_shutdown = True  # No new Function Executor States can be created.
+            while self._function_executor_states:
+                id, state = self._function_executor_states.popitem()
+                # Only ongoing tasks who have a reference to the state already can see it.
+                async with state.lock:
+                    await state.shutdown()
+                    # The task running inside the Function Executor will fail because it's destroyed.
 
 
 def _function_id_with_version(task: Task) -> str:

indexify/function_executor/README.md

@@ -0,0 +1,18 @@
+## Overview
+
+Function Executor is a process with an API that allows to load and run a customer Function in Indexify.
+Each function run is a task. The tasks can be executed concurrently. The API client controls
+the desired concurrency. Killing the process allows to free all the resources that a loaded customer
+functon is using. This is helpful because the SDK doesn't provide any callbacks to customer code to free
+resources it's using. Even if there was such callback customer code still might misbehave.
+
+## Deployment
+
+A Function Executor is created and destroyed by another component called Executor. It also calls the
+Function Executor APIs. The server is not expected to be deployed or managed manually by Indexify users
+as it's a low level component.
+
+## Threat model
+
+Customer code is assumed to be not trusted. Function Executor must not obtain any credentials that grant
+access to resources not owned by the customer who owns the function.

indexify/function_executor/handlers/run_function/function_inputs_loader.py

@@ -1,18 +1,17 @@
 from typing import Optional
 
-from pydantic import BaseModel
+from tensorlake.functions_sdk.data_objects import TensorlakeData
+from tensorlake.functions_sdk.object_serializer import get_serializer
 
-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
+from ...proto.function_executor_pb2 import RunTaskRequest, SerializedObject
 
 
-class FunctionInputs(BaseModel):
-    input: IndexifyData
-    init_value: Optional[IndexifyData] = None
+class FunctionInputs:
+    def __init__(
+        self, input: TensorlakeData, init_value: Optional[TensorlakeData] = None
+    ):
+        self.input = input
+        self.init_value = init_value
 
 
 class FunctionInputsLoader:
@@ -25,12 +24,12 @@ class FunctionInputsLoader:
             init_value=self._accumulator_input(),
         )
 
-    def _function_input(self) -> IndexifyData:
+    def _function_input(self) -> TensorlakeData:
         return _to_indexify_data(
             self._request.graph_invocation_id, self._request.function_input
         )
 
-    def _accumulator_input(self) -> Optional[IndexifyData]:
+    def _accumulator_input(self) -> Optional[TensorlakeData]:
         return (
             _to_indexify_data(
                 self._request.graph_invocation_id, self._request.function_init_value
@@ -42,8 +41,8 @@ class FunctionInputsLoader:
 
 def _to_indexify_data(
     input_id: str, serialized_object: SerializedObject
-) -> IndexifyData:
-    return IndexifyData(
+) -> TensorlakeData:
+    return TensorlakeData(
         input_id=input_id,
         payload=(
             serialized_object.bytes

indexify/function_executor/handlers/run_function/handler.py

@@ -2,25 +2,19 @@ import io
 import sys
 import traceback
 from contextlib import redirect_stderr, redirect_stdout
-from typing import Any, Optional, Union
+from typing import Any, Union
 
-from indexify.function_executor.proto.function_executor_pb2 import (
-    RunTaskRequest,
-    RunTaskResponse,
-)
-from indexify.functions_sdk.indexify_functions import (
+from tensorlake.functions_sdk.functions import (
     FunctionCallResult,
     GraphInvocationContext,
-    IndexifyFunction,
-    IndexifyFunctionWrapper,
-    IndexifyRouter,
     RouterCallResult,
+    TensorlakeCompute,
+    TensorlakeFunctionWrapper,
+    TensorlakeRouter,
 )
-from indexify.functions_sdk.invocation_state.invocation_state import (
-    InvocationState,
-)
-from indexify.http_client import IndexifyClient
+from tensorlake.functions_sdk.invocation_state.invocation_state import InvocationState
 
+from ...proto.function_executor_pb2 import RunTaskRequest, RunTaskResponse
 from .function_inputs_loader import FunctionInputs, FunctionInputsLoader
 from .response_helper import ResponseHelper
 
@@ -30,9 +24,9 @@ class Handler:
         self,
         request: RunTaskRequest,
         graph_name: str,
-        graph_version: int,
+        graph_version: str,
         function_name: str,
-        function: Union[IndexifyFunction, IndexifyRouter],
+        function: Union[TensorlakeCompute, TensorlakeCompute],
         invocation_state: InvocationState,
         logger: Any,
     ):
@@ -48,12 +42,12 @@ class Handler:
         self._func_stdout: io.StringIO = io.StringIO()
         self._func_stderr: io.StringIO = io.StringIO()
 
-        self._function_wrapper: IndexifyFunctionWrapper = IndexifyFunctionWrapper(
+        self._function_wrapper: TensorlakeFunctionWrapper = TensorlakeFunctionWrapper(
             indexify_function=function,
             context=GraphInvocationContext(
                 invocation_id=request.graph_invocation_id,
                 graph_name=graph_name,
-                graph_version=str(graph_version),
+                graph_version=graph_version,
                 invocation_state=invocation_state,
             ),
         )
@@ -116,35 +110,17 @@ class Handler:
         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:
+def _is_router(func_wrapper: TensorlakeFunctionWrapper) -> bool:
     """Determines if the function is a router.
 
-    A function is a router if it is an instance of IndexifyRouter or if it is an IndexifyRouter class.
+    A function is a router if it is an instance of TensorlakeRouter or if it is an TensorlakeRouter class.
     """
     return str(
         type(func_wrapper.indexify_function)
-    ) == "<class 'indexify.functions_sdk.indexify_functions.IndexifyRouter'>" or isinstance(
-        func_wrapper.indexify_function, IndexifyRouter
+    ) == "<class 'tensorlake.functions_sdk.functions.TensorlakeRouter'>" or isinstance(
+        func_wrapper.indexify_function, TensorlakeRouter
     )
 
 
-def _func_is_reducer(func_wrapper: IndexifyFunctionWrapper) -> bool:
+def _func_is_reducer(func_wrapper: TensorlakeFunctionWrapper) -> bool:
     return func_wrapper.indexify_function.accumulate is not None

indexify/function_executor/handlers/run_function/request_validator.py

@@ -1,9 +1,7 @@
 from typing import Any
 
-from indexify.function_executor.proto.function_executor_pb2 import (
-    RunTaskRequest,
-)
-from indexify.function_executor.proto.message_validator import MessageValidator
+from ...proto.function_executor_pb2 import RunTaskRequest
+from ...proto.message_validator import MessageValidator
 
 
 class RequestValidator:
@@ -17,7 +15,11 @@ class RequestValidator:
         Raises: ValueError: If the request is invalid.
         """
         (
-            self._message_validator.required_field("graph_invocation_id")
+            self._message_validator.required_field("namespace")
+            .required_field("graph_name")
+            .required_field("graph_version")
+            .required_field("function_name")
+            .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

@@ -1,17 +1,15 @@
 from typing import List
 
-from indexify.function_executor.proto.function_executor_pb2 import (
+from tensorlake.functions_sdk.data_objects import TensorlakeData
+from tensorlake.functions_sdk.functions import FunctionCallResult, RouterCallResult
+from tensorlake.functions_sdk.object_serializer import get_serializer
+
+from ...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:
@@ -81,7 +79,7 @@ class ResponseHelper:
             success=False,
         )
 
-    def _to_function_output(self, outputs: List[IndexifyData]) -> FunctionOutput:
+    def _to_function_output(self, outputs: List[TensorlakeData]) -> FunctionOutput:
         output = FunctionOutput(outputs=[])
         for ix_data in outputs:
             serialized_object: SerializedObject = SerializedObject(

indexify/function_executor/initialize_request_validator.py

@@ -1,6 +1,5 @@
-from indexify.function_executor.proto.message_validator import MessageValidator
-
 from .proto.function_executor_pb2 import InitializeRequest
+from .proto.message_validator import MessageValidator
 
 
 class InitializeRequestValidator:

indexify/function_executor/invocation_state/invocation_state_proxy_server.py

@@ -2,7 +2,7 @@ import queue
 import threading
 from typing import Any, Iterator, Optional
 
-from indexify.functions_sdk.object_serializer import (
+from tensorlake.functions_sdk.object_serializer import (
     CloudPickleSerializer,
     get_serializer,
 )

indexify/function_executor/invocation_state/proxied_invocation_state.py

@@ -1,8 +1,6 @@
 from typing import Any, Optional
 
-from indexify.functions_sdk.invocation_state.invocation_state import (
-    InvocationState,
-)
+from tensorlake.functions_sdk.invocation_state.invocation_state import InvocationState
 
 from .invocation_state_proxy_server import InvocationStateProxyServer
 

indexify/function_executor/main.py

@@ -0,0 +1,50 @@
+from indexify.utils.logging import (
+    configure_development_mode_logging,
+    configure_logging_early,
+    configure_production_mode_logging,
+)
+
+configure_logging_early()
+
+import argparse
+
+import structlog
+
+from .server import Server
+from .service import Service
+
+logger = structlog.get_logger(module=__name__)
+
+
+def validate_args(args):
+    if args.address is None:
+        logger.error("--address argument is required")
+        exit(1)
+
+
+def main():
+    parser = argparse.ArgumentParser(
+        description="Runs Function Executor with the specified API server address"
+    )
+    parser.add_argument("--address", help="API server address to listen on", type=str)
+    parser.add_argument(
+        "-d", "--dev", help="Run in development mode", action="store_true"
+    )
+    args = parser.parse_args()
+
+    if args.dev:
+        configure_development_mode_logging()
+    else:
+        configure_production_mode_logging()
+    validate_args(args)
+
+    logger.info("starting function executor server", address=args.address)
+
+    Server(
+        server_address=args.address,
+        service=Service(),
+    ).run()
+
+
+if __name__ == "__main__":
+    main()

indexify/function_executor/proto/configuration.py

@@ -4,10 +4,18 @@
 # 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
+# Disable port reuse: fail if multiple Function Executor Servers attempt to bind to the
+# same port. This happens when Indexify users misconfigure the Servers. Disabling the port
+# reuse results in a clear error message on Server startup instead of obscure errors later
+# while Indexify cluster is serving tasks.
+# If we don't disable port reuse then a random Server gets the requests so wrong tasks get
+# routed to wrong servers.
+_REUSE_SERVER_PORT = 0
 
 GRPC_SERVER_OPTIONS = [
     ("grpc.max_receive_message_length", _MAX_GRPC_MESSAGE_LENGTH),
     ("grpc.max_send_message_length", _MAX_GRPC_MESSAGE_LENGTH),
+    ("grpc.so_reuseport", _REUSE_SERVER_PORT),
 ]
 
 GRPC_CHANNEL_OPTIONS = GRPC_SERVER_OPTIONS

indexify/function_executor/proto/function_executor.proto

@@ -24,13 +24,14 @@ message SerializedObject {
 message InitializeRequest {
     optional string namespace = 1;
     optional string graph_name = 2;
-    optional int32 graph_version = 3;
+    optional string graph_version = 3;
     optional string function_name = 5;
     optional SerializedObject graph = 7;
 }
 
 message InitializeResponse {
     optional bool success = 1;
+    optional string customer_error = 2;
 }
 
 message SetInvocationStateRequest {
@@ -85,10 +86,14 @@ message RouterOutput {
 }
 
 message RunTaskRequest {
-    optional string graph_invocation_id = 4;
+    optional string namespace = 1;
+    optional string graph_name = 2;
+    optional string graph_version = 3;
+    optional string function_name = 4;
+    optional string graph_invocation_id = 5;
     optional string task_id = 6;
-    optional SerializedObject function_input = 9;
-    optional SerializedObject function_init_value = 10;
+    optional SerializedObject function_input = 7;
+    optional SerializedObject function_init_value = 8;
 }
 
 message RunTaskResponse {