indexify 0.3.13__py3-none-any.whl → 0.3.14__py3-none-any.whl

indexify/executor/function_executor/health_checker.py

@@ -20,16 +20,22 @@ from .server.client_configuration import HEALTH_CHECK_TIMEOUT_SEC
 HEALTH_CHECK_POLL_PERIOD_SEC = 10
 
 
+class HealthCheckResult:
+    def __init__(self, is_healthy: bool, reason: str):
+        self.is_healthy: bool = is_healthy
+        self.reason: str = reason
+
+
 class HealthChecker:
     def __init__(self, stub: FunctionExecutorStub, logger: Any):
         self._stub: FunctionExecutorStub = stub
         self._logger: Any = logger.bind(module=__name__)
         self._health_check_loop_task: Optional[asyncio.Task] = None
-        self._health_check_failed_callback: Optional[Callable[[], Awaitable[None]]] = (
-            None
-        )
+        self._health_check_failed_callback: Optional[
+            Callable[[HealthCheckResult], Awaitable[None]]
+        ] = None
 
-    async def check(self) -> bool:
+    async def check(self) -> HealthCheckResult:
         """Runs the health check once and returns the result.
 
         Does not raise any exceptions."""
@@ -40,17 +46,25 @@ class HealthChecker:
                 )
                 if not response.healthy:
                     metric_failed_health_checks.inc()
-                return response.healthy
-            except AioRpcError:
+                return HealthCheckResult(
+                    is_healthy=response.healthy, reason=response.status_message
+                )
+            except AioRpcError as e:
                 metric_failed_health_checks.inc()
                 # Expected exception when there are problems with communication because e.g. the server is unhealthy.
-                return False
+                return HealthCheckResult(
+                    is_healthy=False,
+                    reason=f"Executor side RPC channel error: {str(e)}",
+                )
             except Exception as e:
                 metric_failed_health_checks.inc()
                 self._logger.warning("Got unexpected exception, ignoring", exc_info=e)
-                return False
+                return HealthCheckResult(
+                    is_healthy=False,
+                    reason=f"Unexpected exception in Executor: {str(e)}",
+                )
 
-    def start(self, callback: Callable[[], Awaitable[None]]) -> None:
+    def start(self, callback: Callable[[HealthCheckResult], Awaitable[None]]) -> None:
         """Starts periodic health checks.
 
         The supplied callback is an async function called in the calling thread's
@@ -81,9 +95,10 @@ class HealthChecker:
 
     async def _health_check_loop(self) -> None:
         while True:
-            if not await self.check():
+            result: HealthCheckResult = await self.check()
+            if not result.is_healthy:
                 break
             await asyncio.sleep(HEALTH_CHECK_POLL_PERIOD_SEC)
 
-        asyncio.create_task(self._health_check_failed_callback())
+        asyncio.create_task(self._health_check_failed_callback(result))
         self._health_check_loop_task = None

indexify/executor/function_executor/server/function_executor_server_factory.py

@@ -14,8 +14,11 @@ class FunctionExecutorServerConfiguration:
     configuration parameters or raise an exception if it can't implement
     them."""
 
-    def __init__(self, executor_id: str, image_uri: Optional[str]):
+    def __init__(
+        self, executor_id: str, function_executor_id: str, image_uri: Optional[str]
+    ):
         self.executor_id: str = executor_id
+        self.function_executor_id: str = function_executor_id
         # Container image URI of the Function Executor Server.
         self.image_uri: Optional[str] = image_uri
 

indexify/executor/function_executor/single_task_runner.py

@@ -14,6 +14,7 @@ from tensorlake.function_executor.proto.function_executor_pb2_grpc import (
 from ..api_objects import Task
 from .function_executor import CustomerError, FunctionExecutor
 from .function_executor_state import FunctionExecutorState
+from .health_checker import HealthChecker, HealthCheckResult
 from .metrics.single_task_runner import (
     metric_function_executor_run_task_rpc_errors,
     metric_function_executor_run_task_rpc_latency,
@@ -69,7 +70,12 @@ class SingleTaskRunner:
                 await self._create_function_executor()
             except CustomerError as e:
                 return TaskOutput(
-                    task=self._task_input.task,
+                    task_id=self._task_input.task.id,
+                    namespace=self._task_input.task.namespace,
+                    graph_name=self._task_input.task.compute_graph,
+                    function_name=self._task_input.task.compute_fn,
+                    graph_version=self._task_input.task.graph_version,
+                    graph_invocation_id=self._task_input.task.invocation_id,
                     stderr=str(e),
                     success=False,
                 )
@@ -88,6 +94,7 @@ class SingleTaskRunner:
         config: FunctionExecutorServerConfiguration = (
             FunctionExecutorServerConfiguration(
                 executor_id=self._executor_id,
+                function_executor_id=self._state.id,
                 image_uri=self._task_input.task.image_uri,
             )
         )
@@ -144,24 +151,32 @@ class SingleTaskRunner:
                 ).run_task(request)
             return _task_output(task=self._task_input.task, response=response)
 
-    async def _health_check_failed_callback(self):
+    async def _health_check_failed_callback(self, result: HealthCheckResult):
         # Function Executor destroy due to the periodic health check failure ensures that
         # a running task RPC stuck in unhealthy Function Executor fails immidiately.
         async with self._state.lock:
             if self._state.function_executor is not None:
-                await self._destroy_function_executor_on_failed_health_check()
+                await self._destroy_function_executor_on_failed_health_check(
+                    result.reason
+                )
 
     async def _destroy_existing_function_executor_if_unhealthy(self):
         self._state.check_locked()
         if self._state.function_executor is None:
             return
-        if await self._state.function_executor.health_checker().check():
+        result: HealthCheckResult = (
+            await self._state.function_executor.health_checker().check()
+        )
+        if result.is_healthy:
             return
-        await self._destroy_function_executor_on_failed_health_check()
+        await self._destroy_function_executor_on_failed_health_check(result.reason)
 
-    async def _destroy_function_executor_on_failed_health_check(self):
+    async def _destroy_function_executor_on_failed_health_check(self, reason: str):
         self._state.check_locked()
-        self._logger.error("Health check failed, destroying FunctionExecutor.")
+        self._logger.error(
+            "Function Executor health check failed, destroying Function Executor",
+            health_check_fail_reason=reason,
+        )
         self._state.health_check_failed = True
         await self._state.destroy_function_executor()
 
@@ -220,7 +235,12 @@ def _task_output(task: Task, response: RunTaskResponse) -> TaskOutput:
             raise ValueError(f"Response is missing required field: {field}")
 
     output = TaskOutput(
-        task=task,
+        task_id=task.id,
+        namespace=task.namespace,
+        graph_name=task.compute_graph,
+        function_name=task.compute_fn,
+        graph_version=task.graph_version,
+        graph_invocation_id=task.invocation_id,
         stdout=response.stdout,
         stderr=response.stderr,
         reducer=response.is_reducer,

indexify/executor/function_executor/task_output.py

@@ -13,7 +13,12 @@ class TaskOutput:
 
     def __init__(
         self,
-        task: Task,
+        task_id: str,
+        namespace: str,
+        graph_name: str,
+        function_name: str,
+        graph_version: str,
+        graph_invocation_id: str,
         function_output: Optional[FunctionOutput] = None,
         router_output: Optional[RouterOutput] = None,
         stdout: Optional[str] = None,
@@ -22,7 +27,12 @@ class TaskOutput:
         success: bool = False,
         is_internal_error: bool = False,
     ):
-        self.task = task
+        self.task_id = task_id
+        self.namespace = namespace
+        self.graph_name = graph_name
+        self.function_name = function_name
+        self.graph_version = graph_version
+        self.graph_invocation_id = graph_invocation_id
         self.function_output = function_output
         self.router_output = router_output
         self.stdout = stdout
@@ -32,11 +42,24 @@ class TaskOutput:
         self.is_internal_error = is_internal_error
 
     @classmethod
-    def internal_error(cls, task: Task) -> "TaskOutput":
+    def internal_error(
+        cls,
+        task_id: str,
+        namespace: str,
+        graph_name: str,
+        function_name: str,
+        graph_version: str,
+        graph_invocation_id: str,
+    ) -> "TaskOutput":
         """Creates a TaskOutput for an internal error."""
         # We are not sharing internal error messages with the customer.
         return TaskOutput(
-            task=task,
+            task_id=task_id,
+            namespace=namespace,
+            graph_name=graph_name,
+            function_name=function_name,
+            graph_version=graph_version,
+            graph_invocation_id=graph_invocation_id,
             stderr="Platform failed to execute the function.",
             is_internal_error=True,
         )

indexify/executor/state_reconciler.py

@@ -0,0 +1,288 @@
+import asyncio
+from typing import Any, AsyncGenerator, List, Optional, Set
+
+import grpc
+from tensorlake.function_executor.proto.function_executor_pb2 import (
+    InitializeRequest,
+    SerializedObject,
+)
+
+from indexify.task_scheduler.proto.task_scheduler_pb2 import (
+    DesiredExecutorState,
+    FunctionExecutorDescription,
+    FunctionExecutorStatus,
+    GetDesiredExecutorStatesRequest,
+)
+from indexify.task_scheduler.proto.task_scheduler_pb2_grpc import (
+    TaskSchedulerServiceStub,
+)
+
+from .downloader import Downloader
+from .function_executor.function_executor import CustomerError, FunctionExecutor
+from .function_executor.function_executor_state import FunctionExecutorState
+from .function_executor.function_executor_states_container import (
+    FunctionExecutorStatesContainer,
+)
+from .function_executor.server.function_executor_server_factory import (
+    FunctionExecutorServerConfiguration,
+    FunctionExecutorServerFactory,
+)
+from .function_executor.task_input import TaskInput
+from .function_executor.task_output import TaskOutput
+from .metrics.executor import (
+    METRIC_TASKS_COMPLETED_OUTCOME_ALL,
+    METRIC_TASKS_COMPLETED_OUTCOME_ERROR_CUSTOMER_CODE,
+    METRIC_TASKS_COMPLETED_OUTCOME_ERROR_PLATFORM,
+    METRIC_TASKS_COMPLETED_OUTCOME_SUCCESS,
+    metric_task_completion_latency,
+    metric_task_outcome_report_latency,
+    metric_task_outcome_report_retries,
+    metric_task_outcome_reports,
+    metric_tasks_completed,
+    metric_tasks_fetched,
+    metric_tasks_reporting_outcome,
+)
+from .task_reporter import TaskReporter
+
+
+class ExecutorStateReconciler:
+    def __init__(
+        self,
+        executor_id: str,
+        function_executor_server_factory: FunctionExecutorServerFactory,
+        base_url: str,
+        function_executor_states: FunctionExecutorStatesContainer,
+        config_path: Optional[str],
+        downloader: Downloader,
+        task_reporter: TaskReporter,
+        server_channel: grpc.aio.Channel,
+        logger: Any,
+    ):
+        self._executor_id: str = executor_id
+        self._factory: FunctionExecutorServerFactory = function_executor_server_factory
+        self._base_url: str = base_url
+        self._config_path: Optional[str] = config_path
+        self._downloader: Downloader = downloader
+        self._task_reporter: TaskReporter = task_reporter
+        self._function_executor_states: FunctionExecutorStatesContainer = (
+            function_executor_states
+        )
+        self._stub: TaskSchedulerServiceStub = TaskSchedulerServiceStub(server_channel)
+        self._logger: Any = logger.bind(module=__name__)
+        self._is_shutdown: bool = False
+        self._reconciliation_lock: asyncio.Lock = asyncio.Lock()
+        self._server_last_clock: Optional[int] = None
+
+    async def run(self):
+        desired_states: AsyncGenerator[DesiredExecutorState, None] = (
+            self._stub.get_desired_executor_states(
+                GetDesiredExecutorStatesRequest(executor_id=self._executor_id)
+            )
+        )
+        async for new_state in desired_states:
+            if self._is_shutdown:
+                return
+            new_state: DesiredExecutorState
+            if self._server_last_clock is not None:
+                if self._server_last_clock >= new_state.clock:
+                    continue  # Duplicate or outdated message state sent by Server.
+
+            self._server_last_clock = new_state.clock
+            asyncio.create_task(self._reconcile_state(new_state))
+
+    async def _reconcile_state(self, new_state: DesiredExecutorState):
+        if self._is_shutdown:
+            return
+
+        # Simple non concurrent implementation for now for the PoC.
+        # Obtain this lock to force only a single coroutine doing the reconciliation.
+        async with self._reconciliation_lock:
+            await self._reconcile_function_executors(new_state)
+            # TODO
+            # await self._reconcile_task_allocations(new_state)
+
+    async def shutdown(self):
+        """Shuts down the state reconciler.
+
+        Never raises any exceptions.
+        """
+        self._is_shutdown = True
+
+    async def _reconcile_function_executors(self, desired_state: DesiredExecutorState):
+        desired_function_executor_ids: Set[str] = set()
+        for desired_function_executor in desired_state.function_executors:
+            desired_function_executor: FunctionExecutorDescription
+            desired_function_executor_ids.add(desired_function_executor.id)
+
+            function_executor_state: FunctionExecutorState = (
+                self._function_executor_states.get_or_create_state(
+                    id=desired_function_executor.id,
+                    namespace=desired_function_executor.namespace,
+                    graph_name=desired_function_executor.graph_name,
+                    graph_version=desired_function_executor.graph_version,
+                    function_name=desired_function_executor.function_name,
+                )
+            )
+
+            async with function_executor_state.lock:
+                if (
+                    function_executor_state.status
+                    == FunctionExecutorStatus.FUNCTION_EXECUTOR_STATUS_STOPPED
+                ):
+                    function_executor_state.status = (
+                        FunctionExecutorStatus.FUNCTION_EXECUTOR_STATUS_STARTING_UP
+                    )
+                    try:
+                        function_executor_state.function_executor = (
+                            await self._create_function_executor()
+                        )
+                        function_executor_state.status = (
+                            FunctionExecutorStatus.FUNCTION_EXECUTOR_STATUS_IDLE
+                        )
+                    except CustomerError as e:
+                        function_executor_state.status = (
+                            FunctionExecutorStatus.FUNCTION_EXECUTOR_STATUS_STARTUP_FAILED_CUSTOMER_ERROR
+                        )
+                    except Exception as e:
+                        function_executor_state.status = (
+                            FunctionExecutorStatus.FUNCTION_EXECUTOR_STATUS_STARTUP_FAILED_PLATFORM_ERROR
+                        )
+                        self._logger.error(
+                            f"Failed to create Function Executor", exc_info=e
+                        )
+
+        function_executor_state_ids_to_destroy: List[str] = []
+        async for function_executor_state in self._function_executor_states:
+            function_executor_state: FunctionExecutorState
+            if function_executor_state.id not in desired_function_executor_ids:
+                function_executor_state_ids_to_destroy.append(
+                    function_executor_state.id
+                )
+
+        for function_executor_state_id in function_executor_state_ids_to_destroy:
+            function_executor_state: FunctionExecutorState = (
+                self._function_executor_states.pop_state(function_executor_state_id)
+            )
+            async with function_executor_state.lock:
+                logger = self._function_executor_logger(
+                    id=function_executor_state.id,
+                    namespace=function_executor_state.namespace,
+                    graph_name=function_executor_state.graph_name,
+                    graph_version=function_executor_state.graph_version,
+                    function_name=function_executor_state.function_name,
+                )
+                if (
+                    function_executor_state.status
+                    == FunctionExecutorStatus.FUNCTION_EXECUTOR_STATUS_RUNNING_TASK
+                ):
+                    logger.warning(
+                        "Destroying Function Executor that is running a task. No task output will be reported as this is expected by the Server."
+                    )
+                function_executor_state.status = (
+                    FunctionExecutorStatus.FUNCTION_EXECUTOR_STATUS_STOPPING
+                )
+                await function_executor_state.destroy_function_executor()
+                function_executor_state.status = (
+                    FunctionExecutorStatus.FUNCTION_EXECUTOR_STATUS_STOPPED
+                )
+
+    async def _create_function_executor(
+        self, description: FunctionExecutorDescription
+    ) -> FunctionExecutor:
+        logger = self._function_executor_logger(
+            id=description.id,
+            namespace=description.namespace,
+            graph_name=description.graph_name,
+            graph_version=description.graph_version,
+            function_name=description.function_name,
+        )
+        graph: SerializedObject = await self._downloader.download_graph(
+            namespace=description.namespace,
+            graph_name=description.graph_name,
+            graph_version=description.graph_version,
+            logger=logger,
+        )
+        function_executor: FunctionExecutor = FunctionExecutor(
+            server_factory=self._factory, logger=logger
+        )
+        config: FunctionExecutorServerConfiguration = (
+            FunctionExecutorServerConfiguration(
+                executor_id=self._executor_id,
+                function_executor_id=description.id,
+                image_uri=description.image_uri,
+            )
+        )
+        initialize_request: InitializeRequest = InitializeRequest(
+            namespace=description.namespace,
+            graph_name=description.graph_name,
+            graph_version=description.graph_version,
+            function_name=description.function_name,
+            graph=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:
+            await function_executor.destroy()
+            raise
+
+    async def _cancel_running_tasks(
+        self, function_executor_state: FunctionExecutorState
+    ):
+        pass
+
+    def _function_executor_logger(
+        self,
+        id: str,
+        namespace: str,
+        graph_name: str,
+        graph_version: str,
+        function_name: str,
+    ) -> Any:
+        return self._logger.bind(
+            id=id,
+            namespace=namespace,
+            graph=graph_name,
+            graph_version=graph_version,
+            function_name=function_name,
+        )
+
+    async def _report_task_outcome(self, task_output: TaskOutput):
+        """Reports the task with the given output to the server.
+
+        Doesn't raise any Exceptions. Runs till the reporting is successful."""
+        reporting_retries: int = 0
+
+        while True:
+            logger = logger.bind(retries=reporting_retries)
+            try:
+                await self._task_reporter.report(output=task_output, logger=logger)
+                break
+            except Exception as e:
+                logger.error(
+                    "failed to report task",
+                    exc_info=e,
+                )
+                reporting_retries += 1
+                metric_task_outcome_report_retries.inc()
+                await asyncio.sleep(5)
+
+        metric_tasks_completed.labels(outcome=METRIC_TASKS_COMPLETED_OUTCOME_ALL).inc()
+        if task_output.is_internal_error:
+            metric_tasks_completed.labels(
+                outcome=METRIC_TASKS_COMPLETED_OUTCOME_ERROR_PLATFORM
+            ).inc()
+        elif task_output.success:
+            metric_tasks_completed.labels(
+                outcome=METRIC_TASKS_COMPLETED_OUTCOME_SUCCESS
+            ).inc()
+        else:
+            metric_tasks_completed.labels(
+                outcome=METRIC_TASKS_COMPLETED_OUTCOME_ERROR_CUSTOMER_CODE
+            ).inc()

indexify/executor/state_reporter.py

@@ -0,0 +1,127 @@
+import asyncio
+from typing import Any, List, Optional
+
+import grpc
+
+from indexify.task_scheduler.proto.task_scheduler_pb2 import (
+    AllowedFunction,
+    ExecutorState,
+    ExecutorStatus,
+    FunctionExecutorDescription,
+)
+from indexify.task_scheduler.proto.task_scheduler_pb2 import (
+    FunctionExecutorState as FunctionExecutorStateProto,
+)
+from indexify.task_scheduler.proto.task_scheduler_pb2 import (
+    GPUResources,
+    HostResources,
+    ReportExecutorStateRequest,
+)
+from indexify.task_scheduler.proto.task_scheduler_pb2_grpc import (
+    TaskSchedulerServiceStub,
+)
+
+from .api_objects import FunctionURI
+from .function_executor.function_executor_state import FunctionExecutorState
+from .function_executor.function_executor_states_container import (
+    FunctionExecutorStatesContainer,
+)
+
+_REPORTING_INTERVAL_SEC = 5
+_REPORT_RPC_TIMEOUT_SEC = 5
+
+
+class ExecutorStateReporter:
+    def __init__(
+        self,
+        executor_id: str,
+        function_allowlist: Optional[List[FunctionURI]],
+        function_executor_states: FunctionExecutorStatesContainer,
+        server_channel: grpc.aio.Channel,
+        logger: Any,
+    ):
+        self._executor_id: str = executor_id
+        self._function_executor_states: FunctionExecutorStatesContainer = (
+            function_executor_states
+        )
+        self._stub: TaskSchedulerServiceStub = TaskSchedulerServiceStub(server_channel)
+        self._logger: Any = logger.bind(module=__name__)
+        self._is_shutdown: bool = False
+        self._executor_status: ExecutorStatus = ExecutorStatus.EXECUTOR_STATUS_UNKNOWN
+        self._allowed_functions: List[AllowedFunction] = []
+
+        for function_uri in (
+            function_allowlist if function_allowlist is not None else []
+        ):
+            allowed_function = AllowedFunction(
+                namespace=function_uri.namespace,
+                graph_name=function_uri.compute_graph,
+                function_name=function_uri.compute_fn,
+            )
+            if function_uri.version is not None:
+                allowed_function.graph_version = function_uri.version
+            self._allowed_functions.append(allowed_function)
+
+    # TODO: Update Executor to call status updates.
+    def update_status(self, value: ExecutorStatus):
+        self._executor_status = value
+
+    async def run(self):
+        while not self._is_shutdown:
+            await self._report_state()
+            await asyncio.sleep(_REPORTING_INTERVAL_SEC)
+
+    async def _report_state(self):
+        state = ExecutorState(
+            executor_id=self._executor_id,
+            executor_status=self._executor_status,
+            host_resources=await self._fetch_host_resources(),
+            allowed_functions=self._allowed_functions,
+            function_executor_states=await self._fetch_function_executor_states(),
+        )
+
+        await self._stub.report_executor_state(
+            ReportExecutorStateRequest(executor_state=state),
+            timeout=_REPORT_RPC_TIMEOUT_SEC,
+        )
+
+    async def _fetch_host_resources(self) -> HostResources:
+        # We're only supporting Executors with non empty function allowlist right now.
+        # In this mode Server should ignore available host resources.
+        # This is why it's okay to report zeros right now.
+        return HostResources(
+            cpu_count=0,
+            memory_bytes=0,
+            disk_bytes=0,
+            gpu=GPUResources(
+                count=0,
+                model="",
+            ),
+        )
+
+    async def _fetch_function_executor_states(self) -> List[FunctionExecutorStateProto]:
+        states = []
+
+        async for function_executor_state in self._function_executor_states:
+            function_executor_state: FunctionExecutorState
+            states.append(
+                FunctionExecutorStateProto(
+                    description=FunctionExecutorDescription(
+                        id=function_executor_state.id,
+                        namespace=function_executor_state.namespace,
+                        graph_name=function_executor_state.graph_name,
+                        graph_version=function_executor_state.graph_version,
+                        function_name=function_executor_state.function_name,
+                    ),
+                    status=function_executor_state.status,
+                )
+            )
+
+        return states
+
+    async def shutdown(self):
+        """Shuts down the state reporter.
+
+        Never raises any exceptions.
+        """
+        self._is_shutdown = True

indexify/executor/task_reporter.py

@@ -13,12 +13,12 @@ from .api_objects import (
     RouterOutput,
     TaskResult,
 )
+from .function_executor.task_output import TaskOutput
 from .metrics.task_reporter import (
     metric_server_ingest_files_errors,
     metric_server_ingest_files_latency,
     metric_server_ingest_files_requests,
 )
-from .task_runner import TaskOutput
 
 
 # https://github.com/psf/requests/issues/1081#issuecomment-428504128
@@ -118,12 +118,12 @@ class TaskReporter:
     ) -> Tuple[TaskResult, List[Any], TaskOutputSummary]:
         task_result = TaskResult(
             outcome="failure",
-            namespace=output.task.namespace,
-            compute_graph=output.task.compute_graph,
-            compute_fn=output.task.compute_fn,
-            invocation_id=output.task.invocation_id,
+            namespace=output.namespace,
+            compute_graph=output.graph_name,
+            compute_fn=output.function_name,
+            invocation_id=output.graph_invocation_id,
             executor_id=self._executor_id,
-            task_id=output.task.id,
+            task_id=output.task_id,
         )
         output_files: List[Any] = []
         summary: TaskOutputSummary = TaskOutputSummary()