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

indexify/executor/function_executor/process_function_executor_factory.py

@@ -0,0 +1,102 @@
+import asyncio
+from typing import Any, Optional
+
+from .function_executor_factory import FunctionExecutorFactory
+from .process_function_executor import ProcessFunctionExecutor
+
+
+class ProcessFunctionExecutorFactory(FunctionExecutorFactory):
+    def __init__(
+        self,
+        indexify_server_address: str,
+        development_mode: bool,
+        config_path: Optional[str],
+    ):
+        self._indexify_server_address: str = indexify_server_address
+        self._development_mode: bool = development_mode
+        self._config_path: Optional[str] = config_path
+        # Registred ports range end at 49151. We start from 50000 to hopefully avoid conflicts.
+        self._free_ports = set(range(50000, 51000))
+
+    async def create(
+        self, logger: Any, state: Optional[Any] = None
+    ) -> ProcessFunctionExecutor:
+        logger = logger.bind(module=__name__)
+        port: Optional[int] = None
+
+        try:
+            port = self._allocate_port()
+            args = [
+                "function-executor",
+                "--function-executor-server-address",
+                _server_address(port),
+                "--indexify-server-address",
+                self._indexify_server_address,
+            ]
+            if self._development_mode:
+                args.append("--dev")
+            if self._config_path is not None:
+                args.extend(["--config-path", self._config_path])
+            # Run the process with our stdout, stderr. We want to see process logs and exceptions in our process output.
+            # This is useful for dubugging. Customer function stdout and stderr is captured and returned in the response
+            # so we won't see it in our process outputs. This is the right behavior as customer function stdout and stderr
+            # contains private customer data.
+            proc: asyncio.subprocess.Process = await asyncio.create_subprocess_exec(
+                "indexify-cli",
+                *args,
+            )
+            return ProcessFunctionExecutor(
+                process=proc,
+                port=port,
+                address=_server_address(port),
+                logger=logger,
+                state=state,
+            )
+        except Exception as e:
+            if port is not None:
+                self._release_port(port)
+            logger.error(
+                "failed starting a new Function Executor process at port {port}",
+                exc_info=e,
+            )
+            raise
+
+    async def destroy(self, executor: ProcessFunctionExecutor, logger: Any) -> None:
+        proc: asyncio.subprocess.Process = executor._proc
+        port: int = executor._port
+        logger = logger.bind(
+            module=__name__,
+            pid=proc.pid,
+            port=port,
+        )
+
+        try:
+            if proc.returncode is not None:
+                # The process already exited and was waited() sucessfully.
+                return
+
+            proc.kill()
+            await proc.wait()
+        except Exception as e:
+            logger.error(
+                "failed to cleanup Function Executor process",
+                exc_info=e,
+            )
+        finally:
+            self._release_port(port)
+            if executor._channel is not None:
+                await executor._channel.close()
+
+    def _allocate_port(self) -> int:
+        # No asyncio.Lock is required here because this operation never awaits
+        # and it is always called from the same thread where the event loop is running.
+        return self._free_ports.pop()
+
+    def _release_port(self, port: int) -> None:
+        # No asyncio.Lock is required here because this operation never awaits
+        # and it is always called from the same thread where the event loop is running.
+        self._free_ports.add(port)
+
+
+def _server_address(port: int) -> str:
+    return f"localhost:{port}"

indexify/executor/function_worker.py

@@ -1,212 +1,255 @@
-import sys
-import traceback
-from typing import Dict, List, Optional
-
-import cloudpickle
-from pydantic import BaseModel
-from rich import print
-
-from indexify import IndexifyClient
-from indexify.functions_sdk.data_objects import (
-    FunctionWorkerOutput,
-    IndexifyData,
+import asyncio
+from typing import Any, Dict, Optional
+
+import grpc
+import structlog
+
+from indexify.function_executor.proto.function_executor_pb2 import (
+    FunctionOutput,
+    InitializeRequest,
     RouterOutput,
+    RunTaskRequest,
+    RunTaskResponse,
+    SerializedObject,
+)
+from indexify.function_executor.proto.function_executor_pb2_grpc import (
+    FunctionExecutorStub,
 )
-from indexify.functions_sdk.indexify_functions import (
-    FunctionCallResult,
-    GraphInvocationContext,
-    IndexifyFunction,
-    IndexifyFunctionWrapper,
-    IndexifyRouter,
-    RouterCallResult,
+
+from .api_objects import Task
+from .downloader import DownloadedInputs
+from .function_executor.function_executor import FunctionExecutor
+from .function_executor.function_executor_factory import (
+    FunctionExecutorFactory,
 )
+from .function_executor.function_executor_map import FunctionExecutorMap
+
 
-function_wrapper_map: Dict[str, IndexifyFunctionWrapper] = {}
+class FunctionWorkerInput:
+    """Task with all the resources required to run it."""
 
-import concurrent.futures
+    def __init__(
+        self,
+        task: Task,
+        graph: Optional[SerializedObject] = None,
+        function_input: Optional[DownloadedInputs] = None,
+    ):
+        self.task = task
+        # Must not be None when running the task.
+        self.graph = graph
+        # Must not be None when running the task.
+        self.function_input = function_input
 
 
-class FunctionRunException(Exception):
+class FunctionWorkerOutput:
     def __init__(
-        self, exception: Exception, stdout: str, stderr: str, is_reducer: bool
+        self,
+        function_output: Optional[FunctionOutput] = None,
+        router_output: Optional[RouterOutput] = None,
+        stdout: Optional[str] = None,
+        stderr: Optional[str] = None,
+        reducer: bool = False,
+        success: bool = False,
     ):
-        super().__init__(str(exception))
-        self.exception = exception
+        self.function_output = function_output
+        self.router_output = router_output
         self.stdout = stdout
         self.stderr = stderr
-        self.is_reducer = is_reducer
-
-
-class FunctionOutput(BaseModel):
-    fn_outputs: Optional[List[IndexifyData]]
-    router_output: Optional[RouterOutput]
-    reducer: bool = False
-    success: bool = True
-    stdout: str = ""
-    stderr: str = ""
-
-
-def _load_function(
-    namespace: str,
-    graph_name: str,
-    fn_name: str,
-    code_path: str,
-    version: int,
-    invocation_id: str,
-    indexify_client: IndexifyClient,
-):
-    """Load an extractor to the memory: extractor_wrapper_map."""
-    global function_wrapper_map
-    key = f"{namespace}/{graph_name}/{version}/{fn_name}"
-    if key in function_wrapper_map:
-        return
-    with open(code_path, "rb") as f:
-        code = f.read()
-        pickled_functions = cloudpickle.loads(code)
-    context = GraphInvocationContext(
-        invocation_id=invocation_id,
-        graph_name=graph_name,
-        graph_version=str(version),
-        indexify_client=indexify_client,
-    )
-    function_wrapper = IndexifyFunctionWrapper(
-        cloudpickle.loads(pickled_functions[fn_name]),
-        context,
-    )
-    function_wrapper_map[key] = function_wrapper
+        self.reducer = reducer
+        self.success = success
 
 
-class FunctionWorker:
+class FunctionExecutorState:
     def __init__(
-        self, workers: int = 1, indexify_client: IndexifyClient = None
-    ) -> None:
-        self._executor: concurrent.futures.ProcessPoolExecutor = (
-            concurrent.futures.ProcessPoolExecutor(max_workers=workers)
-        )
-        self._workers = workers
-        self._indexify_client = indexify_client
-
-    async def async_submit(
         self,
-        namespace: str,
-        graph_name: str,
-        fn_name: str,
-        input: IndexifyData,
-        code_path: str,
-        version: int,
-        init_value: Optional[IndexifyData] = None,
-        invocation_id: Optional[str] = None,
-    ) -> FunctionWorkerOutput:
+        function_id_with_version: str,
+        function_id_without_version: str,
+        ongoing_tasks_count: int,
+    ):
+        self.function_id_with_version: str = function_id_with_version
+        self.function_id_without_version: str = function_id_without_version
+        self.ongoing_tasks_count: int = ongoing_tasks_count
+
+
+class FunctionWorker:
+    def __init__(self, function_executor_factory: FunctionExecutorFactory):
+        self._function_executors = FunctionExecutorMap(function_executor_factory)
+
+    async def run(self, input: FunctionWorkerInput) -> FunctionWorkerOutput:
+        logger = _logger(input.task)
+        function_executor: Optional[FunctionExecutor] = None
         try:
-            result = _run_function(
-                namespace,
-                graph_name,
-                fn_name,
-                input,
-                code_path,
-                version,
-                init_value,
-                invocation_id,
-                self._indexify_client,
+            function_executor = await self._obtain_function_executor(input, logger)
+            return await self._run_in_executor(
+                function_executor=function_executor, input=input
             )
-            # TODO - bring back running in a separate process
         except Exception as e:
-            return FunctionWorkerOutput(
-                stdout=e.stdout,
-                stderr=e.stderr,
-                reducer=e.is_reducer,
-                success=False,
+            logger.error(
+                "failed running the task",
+                exc_info=e,
             )
-
-        return FunctionWorkerOutput(
-            fn_outputs=result.fn_outputs,
-            router_output=result.router_output,
-            stdout=result.stdout,
-            stderr=result.stderr,
-            reducer=result.reducer,
-            success=result.success,
+            if function_executor is not None:
+                # This will fail all the tasks concurrently running in this Function Executor. Not great.
+                await self._function_executors.delete(
+                    id=_function_id_without_version(input.task),
+                    function_executor=function_executor,
+                    logger=logger,
+                )
+            return _internal_error_output()
+
+    async def _obtain_function_executor(
+        self, input: FunctionWorkerInput, logger: Any
+    ) -> FunctionExecutor:
+        # Temporary policy for Function Executors lifecycle:
+        # There can only be a single Function Executor per function.
+        # If a Function Executor already exists for a different function version then wait until
+        # all the tasks finish in the existing Function Executor and then destroy it first.
+        initialize_request: InitializeRequest = InitializeRequest(
+            namespace=input.task.namespace,
+            graph_name=input.task.compute_graph,
+            graph_version=input.task.graph_version,
+            function_name=input.task.compute_fn,
+            graph=input.graph,
+        )
+        initial_function_executor_state: FunctionExecutorState = FunctionExecutorState(
+            function_id_with_version=_function_id_with_version(input.task),
+            function_id_without_version=_function_id_without_version(input.task),
+            ongoing_tasks_count=0,
         )
 
-    def shutdown(self):
-        self._executor.shutdown(wait=True, cancel_futures=True)
-
-
-def _run_function(
-    namespace: str,
-    graph_name: str,
-    fn_name: str,
-    input: IndexifyData,
-    code_path: str,
-    version: int,
-    init_value: Optional[IndexifyData] = None,
-    invocation_id: Optional[str] = None,
-    indexify_client: Optional[IndexifyClient] = None,
-) -> FunctionOutput:
-    import io
-    from contextlib import redirect_stderr, redirect_stdout
-
-    stdout_capture = io.StringIO()
-    stderr_capture = io.StringIO()
-    is_reducer = False
-    router_output = None
-    fn_output = None
-    has_failed = False
-    print(
-        f"[bold] function_worker: [/bold] invoking function {fn_name} in graph {graph_name}"
-    )
-    with redirect_stdout(stdout_capture), redirect_stderr(stderr_capture):
-        try:
-            key = f"{namespace}/{graph_name}/{version}/{fn_name}"
-            if key not in function_wrapper_map:
-                _load_function(
-                    namespace,
-                    graph_name,
-                    fn_name,
-                    code_path,
-                    version,
-                    invocation_id,
-                    indexify_client,
-                )
+        while True:
+            function_executor = await self._function_executors.get_or_create(
+                id=_function_id_without_version(input.task),
+                initialize_request=initialize_request,
+                initial_state=initial_function_executor_state,
+                logger=logger,
+            )
 
-            fn = function_wrapper_map[key]
+            # No need to lock Function Executor state as we are not awaiting.
+            function_executor_state: FunctionExecutorState = function_executor.state()
             if (
-                str(type(fn.indexify_function))
-                == "<class 'indexify.functions_sdk.indexify_functions.IndexifyRouter'>"
+                function_executor_state.function_id_with_version
+                == _function_id_with_version(input.task)
             ):
-                router_call_result: RouterCallResult = fn.invoke_router(fn_name, input)
-                router_output = RouterOutput(edges=router_call_result.edges)
-                if router_call_result.traceback_msg is not None:
-                    print(router_call_result.traceback_msg, file=sys.stderr)
-                    has_failed = True
+                # The existing Function Executor is for the same function version so we can run the task in it.
+                # Increment the ongoing tasks count before awaiting to prevent the Function Executor from being destroyed
+                # by another coroutine.
+                function_executor_state.ongoing_tasks_count += 1
+                return function_executor
+
+            # This loop implements the temporary policy so it's implemented using polling instead of a lower
+            # latency event based mechanism with a higher complexity.
+            if function_executor_state.ongoing_tasks_count == 0:
+                logger.info(
+                    "destroying existing Function Executor for different function version",
+                    function_id=_function_id_with_version(input.task),
+                    executor_function_id=function_executor_state.function_id_with_version,
+                )
+                await self._function_executors.delete(
+                    id=_function_id_without_version(input.task),
+                    function_executor=function_executor,
+                    logger=logger,
+                )
             else:
-                fn_call_result: FunctionCallResult = fn.invoke_fn_ser(
-                    fn_name, input, init_value
+                logger.info(
+                    "waiting for existing Function Executor to finish",
+                    function_id=_function_id_with_version(input.task),
+                    executor_function_id=function_executor_state.function_id_with_version,
+                )
+                await asyncio.sleep(
+                    5
+                )  # Wait for 5 secs before checking if all tasks for the existing Function Executor finished.
+
+    async def _run_in_executor(
+        self, function_executor: FunctionExecutor, input: FunctionWorkerInput
+    ) -> FunctionWorkerOutput:
+        """Runs the task in the Function Executor.
+
+        The Function Executor's ongoing_tasks_count must be incremented before calling this function.
+        """
+        try:
+            run_task_request: RunTaskRequest = RunTaskRequest(
+                graph_invocation_id=input.task.invocation_id,
+                task_id=input.task.id,
+                function_input=input.function_input.input,
+            )
+            if input.function_input.init_value is not None:
+                run_task_request.function_init_value.CopyFrom(
+                    input.function_input.init_value
                 )
-                is_reducer = fn.indexify_function.accumulate is not None
-                fn_output = fn_call_result.ser_outputs
-                if fn_call_result.traceback_msg is not None:
-                    print(fn_call_result.traceback_msg, file=sys.stderr)
-                    has_failed = True
-        except Exception:
-            print(traceback.format_exc(), file=sys.stderr)
-            has_failed = True
-
-    # WARNING - IF THIS FAILS, WE WILL NOT BE ABLE TO RECOVER
-    # ANY LOGS
-    if has_failed:
-        return FunctionOutput(
-            fn_outputs=None,
-            router_output=None,
-            stdout=stdout_capture.getvalue(),
-            stderr=stderr_capture.getvalue(),
-            reducer=is_reducer,
-            success=False,
+            channel: grpc.aio.Channel = await function_executor.channel()
+            run_task_response: RunTaskResponse = await FunctionExecutorStub(
+                channel
+            ).run_task(run_task_request)
+            return _to_output(run_task_response)
+        finally:
+            # If this Function Executor was destroyed then it's not
+            # visible in the map but we still have a reference to it.
+            function_executor.state().ongoing_tasks_count -= 1
+
+    async def shutdown(self) -> None:
+        await self._function_executors.clear(
+            logger=structlog.get_logger(module=__name__, event="shutdown")
         )
-    return FunctionOutput(
-        fn_outputs=fn_output,
-        router_output=router_output,
-        reducer=is_reducer,
-        success=True,
-        stdout=stdout_capture.getvalue(),
-        stderr=stderr_capture.getvalue(),
+
+
+def _to_output(response: RunTaskResponse) -> FunctionWorkerOutput:
+    required_fields = [
+        "stdout",
+        "stderr",
+        "is_reducer",
+        "success",
+    ]
+
+    for field in required_fields:
+        if not response.HasField(field):
+            raise ValueError(f"Response is missing required field: {field}")
+
+    output = FunctionWorkerOutput(
+        stdout=response.stdout,
+        stderr=response.stderr,
+        reducer=response.is_reducer,
+        success=response.success,
+    )
+
+    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 _internal_error_output() -> FunctionWorkerOutput:
+    return FunctionWorkerOutput(
+        function_output=None,
+        router_output=None,
+        stdout=None,
+        # We are not sharing internal error messages with the customer.
+        # If customer code failed then we won't get any exceptions here.
+        # All customer code errors are returned in the gRPC response.
+        stderr="Platform failed to execute the function.",
+        reducer=False,
+        success=False,
+    )
+
+
+def _logger(task: Task) -> Any:
+    return structlog.get_logger(
+        module=__name__,
+        namespace=task.namespace,
+        graph_name=task.compute_graph,
+        graph_version=task.graph_version,
+        function_name=task.compute_fn,
+        graph_invocation_id=task.invocation_id,
+        task_id=task.id,
+        function_id=_function_id_with_version(task),
     )
+
+
+def _function_id_with_version(task: Task) -> str:
+    return f"versioned/{task.namespace}/{task.compute_graph}/{task.graph_version}/{task.compute_fn}"
+
+
+def _function_id_without_version(task: Task) -> str:
+    return f"not_versioned/{task.namespace}/{task.compute_graph}/{task.compute_fn}"

indexify/executor/runtime_probes.py

@@ -6,12 +6,13 @@ from typing import Any, Dict, Tuple
 from pydantic import BaseModel
 
 DEFAULT_EXECUTOR = "tensorlake/indexify-executor-default"
-DEFAULT_VERSION = 1
+# Empty string is used as a default hash which tells the scheduler to accept any hash.
+DEFAULT_HASH = ""
 
 
 class ProbeInfo(BaseModel):
     image_name: str
-    image_version: int
+    image_hash: str
     python_major_version: int
     labels: Dict[str, Any] = {}
     is_default_executor: bool
@@ -20,7 +21,7 @@ class ProbeInfo(BaseModel):
 class RuntimeProbes:
     def __init__(self) -> None:
         self._image_name = self._read_image_name()
-        self._image_version = self._read_image_version()
+        self._image_hash = self._read_image_hash()
         self._os_name = platform.system()
         self._architecture = platform.machine()
         (
@@ -35,12 +36,12 @@ class RuntimeProbes:
                 return file.read().strip()
         return DEFAULT_EXECUTOR
 
-    def _read_image_version(self) -> int:
-        file_path = os.path.expanduser("~/.indexify/image_version")
+    def _read_image_hash(self) -> str:
+        file_path = os.path.expanduser("~/.indexify/image_hash")
         if os.path.exists(file_path):
             with open(file_path, "r") as file:
-                return int(file.read().strip())
-        return DEFAULT_VERSION
+                return file.read().strip()
+        return DEFAULT_HASH
 
     def _get_python_version(self) -> Tuple[int, int]:
         version_info = sys.version_info
@@ -60,7 +61,7 @@ class RuntimeProbes:
 
         return ProbeInfo(
             image_name=self._image_name,
-            image_version=self._image_version,
+            image_hash=self._image_hash,
             python_major_version=self._python_version_major,
             labels=labels,
             is_default_executor=self._is_default_executor(),

indexify/executor/task_fetcher.py

@@ -0,0 +1,80 @@
+import json
+from importlib.metadata import version
+from typing import AsyncGenerator, Optional
+
+import httpx
+import structlog
+from httpx_sse import aconnect_sse
+
+from indexify.common_util import get_httpx_client
+
+from .api_objects import ExecutorMetadata, Task
+from .runtime_probes import ProbeInfo, RuntimeProbes
+
+
+class TaskFetcher:
+    """Registers with Indexify server and fetches tasks from it."""
+
+    def __init__(
+        self,
+        protocol: str,
+        indexify_server_addr: str,
+        executor_id: str,
+        name_alias: Optional[str] = None,
+        image_hash: Optional[int] = None,
+        config_path: Optional[str] = None,
+    ):
+        self._protocol: str = protocol
+        self._indexify_server_addr: str = indexify_server_addr
+        self.config_path = config_path
+        self._logger = structlog.get_logger(module=__name__)
+
+        probe_info: ProbeInfo = RuntimeProbes().probe()
+        self._executor_metadata: ExecutorMetadata = ExecutorMetadata(
+            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),
+            labels=probe_info.labels,
+        )
+
+    async def run(self) -> AsyncGenerator[Task, None]:
+        """Fetches tasks that Indexify server assigned to the Executor.
+
+        Raises an exception if error occurred."""
+        url = f"{self._protocol}://{self._indexify_server_addr}/internal/executors/{self._executor_metadata.id}/tasks"
+
+        self._logger.info(
+            "registering_executor",
+            executor_id=self._executor_metadata.id,
+            url=url,
+            executor_version=self._executor_metadata.executor_version,
+        )
+        async with get_httpx_client(
+            config_path=self.config_path, make_async=True
+        ) as client:
+            async with aconnect_sse(
+                client,
+                "POST",
+                url,
+                json=self._executor_metadata.model_dump(),
+                headers={"Content-Type": "application/json"},
+            ) as event_source:
+                try:
+                    event_source.response.raise_for_status()
+                except Exception as e:
+                    await event_source.response.aread()
+                    raise Exception(
+                        "Failed to register at server. "
+                        f"Response code: {event_source.response.status_code}. "
+                        f"Response text: '{event_source.response.text}'."
+                    ) from e
+
+                self._logger.info(
+                    "executor_registered", executor_id=self._executor_metadata.id
+                )
+                async for sse in event_source.aiter_sse():
+                    task_dicts = json.loads(sse.data)
+                    for task_dict in task_dicts:
+                        yield Task.model_validate(task_dict, strict=False)