indexify 0.3.30__py3-none-any.whl → 0.4.2__py3-none-any.whl

indexify/executor/{grpc/state_reconciler.py → state_reconciler.py}

@@ -1,4 +1,5 @@
 import asyncio
+from pathlib import Path
 from typing import Any, AsyncGenerator, Dict, Iterable, List, Optional, Set
 
 from tensorlake.function_executor.proto.message_validator import MessageValidator
@@ -6,26 +7,23 @@ from tensorlake.function_executor.proto.message_validator import MessageValidato
 from indexify.proto.executor_api_pb2 import (
     DesiredExecutorState,
     FunctionExecutorDescription,
+    FunctionExecutorTerminationReason,
     GetDesiredExecutorStatesRequest,
     TaskAllocation,
 )
 from indexify.proto.executor_api_pb2_grpc import ExecutorAPIStub
 
-from ..downloader import Downloader
-from ..function_executor.function_executor_state import FunctionExecutorState
-from ..function_executor.function_executor_states_container import (
-    FunctionExecutorStatesContainer,
-)
-from ..function_executor.function_executor_status import FunctionExecutorStatus
-from ..function_executor.server.function_executor_server_factory import (
+from .blob_store.blob_store import BLOBStore
+from .channel_manager import ChannelManager
+from .function_executor.server.function_executor_server_factory import (
     FunctionExecutorServerFactory,
 )
-from ..task_reporter import TaskReporter
-from .channel_manager import ChannelManager
 from .function_executor_controller import (
     FunctionExecutorController,
     function_executor_logger,
+    task_logger,
     validate_function_executor_description,
+    validate_task,
 )
 from .metrics.state_reconciler import (
     metric_state_reconciliation_errors,
@@ -33,7 +31,6 @@ from .metrics.state_reconciler import (
     metric_state_reconciliations,
 )
 from .state_reporter import ExecutorStateReporter
-from .task_controller import TaskController, task_logger, validate_task
 
 _RECONCILE_STREAM_BACKOFF_INTERVAL_SEC = 5
 _RECONCILIATION_RETRIES = 3
@@ -45,10 +42,9 @@ class ExecutorStateReconciler:
         executor_id: str,
         function_executor_server_factory: FunctionExecutorServerFactory,
         base_url: str,
-        function_executor_states: FunctionExecutorStatesContainer,
         config_path: Optional[str],
-        downloader: Downloader,
-        task_reporter: TaskReporter,
+        cache_path: Path,
+        blob_store: BLOBStore,
         channel_manager: ChannelManager,
         state_reporter: ExecutorStateReporter,
         logger: Any,
@@ -60,21 +56,17 @@ class ExecutorStateReconciler:
         )
         self._base_url: str = base_url
         self._config_path: Optional[str] = config_path
-        self._downloader: Downloader = downloader
-        self._task_reporter: TaskReporter = task_reporter
+        self._cache_path: Path = cache_path
+        self._blob_store: BLOBStore = blob_store
         self._channel_manager: ChannelManager = channel_manager
         self._state_reporter: ExecutorStateReporter = state_reporter
-        self._reconciliation_loop_task: Optional[asyncio.Task] = None
         self._logger: Any = logger.bind(module=__name__)
         self._server_backoff_interval_sec: int = server_backoff_interval_sec
 
         # Mutable state. Doesn't need lock because we access from async tasks running in the same thread.
-        self._is_shutdown: bool = False
-        self._function_executor_states: FunctionExecutorStatesContainer = (
-            function_executor_states
-        )
+        self._desired_states_reader_task: Optional[asyncio.Task] = None
+        self._reconciliation_loop_task: Optional[asyncio.Task] = None
         self._function_executor_controllers: Dict[str, FunctionExecutorController] = {}
-        self._task_controllers: Dict[str, TaskController] = {}
         self._last_server_clock: Optional[int] = None
 
         self._last_desired_state_lock = asyncio.Lock()
@@ -83,23 +75,79 @@ class ExecutorStateReconciler:
         )
         self._last_desired_state: Optional[DesiredExecutorState] = None
 
-    async def run(self):
+    def run(self):
         """Runs the state reconciler.
 
-        Never raises any exceptions.
+        Never raises any exceptions. Doesn't block.
         """
+        if self._reconciliation_loop_task is not None:
+            self._logger.error(
+                "reconciliation loop task is already running, skipping run call"
+            )
+            return
+
         self._reconciliation_loop_task = asyncio.create_task(
             self._reconciliation_loop(),
             name="state reconciler reconciliation loop",
         )
+        self._desired_states_reader_task = asyncio.create_task(
+            self._desired_states_reader_loop(),
+            name="state reconciler desired states stream reader",
+        )
+
+    async def shutdown(self):
+        """Shuts down the state reconciler.
+
+        Never raises any exceptions.
+        """
+        if self._reconciliation_loop_task is not None:
+            self._reconciliation_loop_task.cancel()
+            try:
+                await self._reconciliation_loop_task
+            except asyncio.CancelledError:
+                # Expected cancellation, nothing to do.
+                pass
+            self._logger.info("reconciliation loop is shutdown")
+
+        if self._desired_states_reader_task is not None:
+            self._desired_states_reader_task.cancel()
+            try:
+                await self._desired_states_reader_task
+            except asyncio.CancelledError:
+                # Expected cancellation, nothing to do.
+                pass
+            self._logger.info("desired states stream reader loop is shutdown")
+
+        # Now all the aio tasks exited so nothing will intervene with our actions from this point.
+        fe_shutdown_tasks: List[asyncio.Task] = []
+        for fe_controller in self._function_executor_controllers.values():
+            fe_shutdown_tasks.append(
+                asyncio.create_task(
+                    fe_controller.shutdown(
+                        termination_reason=FunctionExecutorTerminationReason.FUNCTION_EXECUTOR_TERMINATION_REASON_EXECUTOR_SHUTDOWN
+                    ),
+                    name=f"Shutdown Function Executor {fe_controller.function_executor_id()}",
+                )
+            )
+
+        # Run all the shutdown tasks concurrently and wait for them to complete.
+        for task in fe_shutdown_tasks:
+            await task
+
+        self._function_executor_controllers.clear()
+        self._logger.info("state reconciler is shutdown")
+
+    async def _desired_states_reader_loop(self):
+        """Reads the desired states stream from Server and processes it.
 
-        # TODO: Move this into a new async task and cancel it in shutdown().
-        while not self._is_shutdown:
+        Never raises any exceptions. Get cancelled via aio task cancellation.
+        """
+        while True:
             try:
                 stub = ExecutorAPIStub(await self._channel_manager.get_channel())
                 # Report state once before starting the stream so Server
                 # doesn't use stale state it knew about this Executor in the past.
-                await self._state_reporter.report_state(stub)
+                await self._state_reporter.report_state_and_wait_for_completion()
 
                 desired_states_stream: AsyncGenerator[DesiredExecutorState, None] = (
                     stub.get_desired_executor_states(
@@ -123,9 +171,6 @@ class ExecutorStateReconciler:
         self, desired_states: AsyncGenerator[DesiredExecutorState, None]
     ):
         async for new_state in desired_states:
-            if self._is_shutdown:
-                return
-
             new_state: DesiredExecutorState
             validator: MessageValidator = MessageValidator(new_state)
             try:
@@ -155,28 +200,12 @@ class ExecutorStateReconciler:
                 self._last_desired_state = new_state
                 self._last_desired_state_change_notifier.notify_all()
 
-    async def shutdown(self):
-        """Shuts down the state reconciler.
-
-        Never raises any exceptions.
-        """
-        self._is_shutdown = True
-        if self._reconciliation_loop_task is not None:
-            self._reconciliation_loop_task.cancel()
-            self._logger.info("reconciliation loop shutdown")
-
-        for controller in self._task_controllers.values():
-            await controller.destroy()
-        # FEs are destroyed in executor.py right now.
-        # TODO: Once HTTP loop is removed add all FE state and controllers
-        # shutdown logic here. This should allow us to get rid of hacky
-        # "cancel all tasks loop" in executor.py shutdown and make the shutdown
-        # much more controllable and clean. E.g. we would be able to remove logs
-        # suppression from shutdown logic. Also need to shutdown self._function_executor_controllers.
-
     async def _reconciliation_loop(self):
+        """Continuously reconciles the desired state with the current state.
+
+        Never raises any exceptions. Get cancelled via aio task cancellation."""
         last_reconciled_state: Optional[DesiredExecutorState] = None
-        while not self._is_shutdown:
+        while True:
             async with self._last_desired_state_lock:
                 # Comparing object identities (references) is enough here to not reconcile
                 # the same state twice.
@@ -200,10 +229,8 @@ class ExecutorStateReconciler:
         for attempt in range(_RECONCILIATION_RETRIES):
             try:
                 # Reconcile FEs first because Tasks depend on them.
-                await self._reconcile_function_executors(
-                    desired_state.function_executors
-                )
-                await self._reconcile_tasks(desired_state.task_allocations)
+                self._reconcile_function_executors(desired_state.function_executors)
+                self._reconcile_tasks(desired_state.task_allocations)
                 return
             except Exception as e:
                 self._logger.error(
@@ -219,30 +246,20 @@ class ExecutorStateReconciler:
             f"failed to reconcile desired state after {_RECONCILIATION_RETRIES} attempts",
         )
 
-    async def _reconcile_function_executors(
+    def _reconcile_function_executors(
         self, function_executor_descriptions: Iterable[FunctionExecutorDescription]
     ):
         valid_fe_descriptions: List[FunctionExecutorDescription] = (
             self._valid_function_executor_descriptions(function_executor_descriptions)
         )
         for fe_description in valid_fe_descriptions:
-            await self._reconcile_function_executor(fe_description)
+            self._reconcile_function_executor(fe_description)
 
         seen_fe_ids: Set[str] = set(map(lambda fe: fe.id, valid_fe_descriptions))
         fe_ids_to_remove = set(self._function_executor_controllers.keys()) - seen_fe_ids
-        for function_executor_id in fe_ids_to_remove:
-            # Remove the controller before FE shutdown completes so we won't attempt to do it
-            # again on the next reconciliations.
-            await self._function_executor_controllers.pop(
-                function_executor_id
-            ).shutdown()
-            # Schedule removal of the FE state after shutdown. This is required for Server
-            # to known when exactly FE resources are freed so it can put a replacement FE if needed.
-            # Running in a separate asyncio task because this will block until the shutdown is complete.
-            asyncio.create_task(
-                self._remove_function_executor_after_shutdown(function_executor_id),
-                name="Remove Function Executor after shutdown",
-            )
+        for fe_id in fe_ids_to_remove:
+            # Server forgot this FE, so its safe to forget about it now too.
+            self._remove_function_executor_controller(fe_id)
 
     def _valid_function_executor_descriptions(
         self, function_executor_descriptions: Iterable[FunctionExecutorDescription]
@@ -267,7 +284,7 @@ class ExecutorStateReconciler:
 
         return valid_function_executor_descriptions
 
-    async def _reconcile_function_executor(
+    def _reconcile_function_executor(
         self, function_executor_description: FunctionExecutorDescription
     ):
         """Reconciles a single Function Executor with the desired state.
@@ -275,114 +292,94 @@ class ExecutorStateReconciler:
         Doesn't block on any long running operations. Doesn't raise any exceptions.
         """
 
-        if not self._function_executor_states.exists(function_executor_description.id):
-            await self._create_function_executor(function_executor_description)
+        if function_executor_description.id not in self._function_executor_controllers:
+            self._add_function_executor_controller(function_executor_description)
 
-    async def _create_function_executor(
+    def _add_function_executor_controller(
         self, function_executor_description: FunctionExecutorDescription
     ) -> None:
-        """Creates Function Executor for the supplied description.
+        """Creates Function Executor for the supplied description and adds it to internal data structures.
 
         Doesn't block on any long running operations. Doesn't raise any exceptions.
         """
         logger = function_executor_logger(function_executor_description, self._logger)
         try:
-            # TODO: Store FE description in FE state object once we migrate to gRPC State Reconciler.
-            # Then most of these parameters will be removed. Also remove the container and use a simple
-            # Dict once FE shutdown logic is moved into reconciler.
-            function_executor_state: FunctionExecutorState = (
-                await self._function_executor_states.get_or_create_state(
-                    id=function_executor_description.id,
-                    namespace=function_executor_description.namespace,
-                    graph_name=function_executor_description.graph_name,
-                    graph_version=function_executor_description.graph_version,
-                    function_name=function_executor_description.function_name,
-                    image_uri=(
-                        function_executor_description.image_uri
-                        if function_executor_description.HasField("image_uri")
-                        else None
-                    ),
-                    secret_names=list(function_executor_description.secret_names),
-                )
-            )
             controller: FunctionExecutorController = FunctionExecutorController(
                 executor_id=self._executor_id,
-                function_executor_state=function_executor_state,
                 function_executor_description=function_executor_description,
                 function_executor_server_factory=self._function_executor_server_factory,
-                downloader=self._downloader,
+                state_reporter=self._state_reporter,
+                blob_store=self._blob_store,
                 base_url=self._base_url,
                 config_path=self._config_path,
+                cache_path=self._cache_path,
                 logger=self._logger,
             )
             self._function_executor_controllers[function_executor_description.id] = (
                 controller
             )
-            # Ask the controller to create the new FE. Task controllers will notice that the FE is eventually
-            # IDLE and start running tasks on it. Server currently doesn't explicitly manage the desired FE status.
-            await controller.startup()
+            controller.startup()
         except Exception as e:
             logger.error("failed adding Function Executor", exc_info=e)
 
-    async def _remove_function_executor_after_shutdown(
-        self, function_executor_id: str
-    ) -> None:
-        fe_state: FunctionExecutorState = await self._function_executor_states.get(
-            function_executor_id
+    def _remove_function_executor_controller(self, function_executor_id: str) -> None:
+        fe_controller: FunctionExecutorController = (
+            self._function_executor_controllers.pop(function_executor_id)
+        )
+        asyncio.create_task(
+            fe_controller.shutdown(
+                termination_reason=FunctionExecutorTerminationReason.FUNCTION_EXECUTOR_TERMINATION_REASON_REMOVED_FROM_DESIRED_STATE
+            ),
+            name=f"Shutdown Function Executor {function_executor_id}",
         )
-        async with fe_state.lock:
-            await fe_state.wait_status(allowlist=[FunctionExecutorStatus.SHUTDOWN])
-        # The whole reconciler could shutdown while we were waiting for the FE to shutdown.
-        if not self._is_shutdown:
-            await self._function_executor_states.pop(function_executor_id)
 
-    async def _reconcile_tasks(self, task_allocations: Iterable[TaskAllocation]):
+    def _reconcile_tasks(self, task_allocations: Iterable[TaskAllocation]):
         valid_task_allocations: List[TaskAllocation] = self._valid_task_allocations(
             task_allocations
         )
         for task_allocation in valid_task_allocations:
-            await self._reconcile_task(task_allocation)
+            self._reconcile_task(task_allocation)
 
-        seen_task_ids: Set[str] = set(
-            map(lambda task_allocation: task_allocation.task.id, valid_task_allocations)
-        )
-        task_ids_to_remove = set(self._task_controllers.keys()) - seen_task_ids
-        for task_id in task_ids_to_remove:
-            await self._remove_task(task_id)
+        # Cancel tasks that are no longer in the desired state.
+        # FE ID => [Task ID]
+        desired_task_ids_per_fe: Dict[str, List[str]] = {}
+        for task_allocation in valid_task_allocations:
+            if task_allocation.function_executor_id not in desired_task_ids_per_fe:
+                desired_task_ids_per_fe[task_allocation.function_executor_id] = []
+            desired_task_ids_per_fe[task_allocation.function_executor_id].append(
+                task_allocation.task.id
+            )
 
-    async def _reconcile_task(self, task_allocation: TaskAllocation):
+        for fe_controller in self._function_executor_controllers.values():
+            fe_controller: FunctionExecutorController
+            if fe_controller.function_executor_id() in desired_task_ids_per_fe:
+                desired_fe_task_ids: Set[str] = set(
+                    desired_task_ids_per_fe[fe_controller.function_executor_id()]
+                )
+            else:
+                # No tasks desired for this FE, so cancel all its tasks.
+                desired_fe_task_ids: Set[str] = set()
+            actual_fe_task_ids: Set[str] = set(fe_controller.task_ids())
+            task_ids_to_remove: Set[str] = actual_fe_task_ids - desired_fe_task_ids
+            for task_id in task_ids_to_remove:
+                fe_controller.remove_task(task_id)
+
+    def _reconcile_task(self, task_allocation: TaskAllocation):
         """Reconciles a single TaskAllocation with the desired state.
 
         Doesn't raise any exceptions.
         """
-        if task_allocation.task.id in self._task_controllers:
-            # Nothing to do, task allocation already exists and it's immutable.
+        function_executor_controller: FunctionExecutorController = (
+            self._function_executor_controllers[task_allocation.function_executor_id]
+        )
+        if function_executor_controller.has_task(task_allocation.task.id):
+            # Nothing to do, task already exists and it's immutable.
             return
 
-        logger = self._task_allocation_logger(task_allocation)
-        try:
-            function_executor_state: FunctionExecutorState = (
-                await self._function_executor_states.get(
-                    task_allocation.function_executor_id
-                )
-            )
-            self._task_controllers[task_allocation.task.id] = TaskController(
-                task=task_allocation.task,
-                downloader=self._downloader,
-                task_reporter=self._task_reporter,
-                function_executor_id=task_allocation.function_executor_id,
-                function_executor_state=function_executor_state,
-                logger=self._logger,
-            )
-        except Exception as e:
-            logger.error("failed adding TaskController", exc_info=e)
-
-    async def _remove_task(self, task_id: str) -> None:
-        """Schedules removal of an existing task.
-
-        Doesn't block on any long running operations. Doesn't raise any exceptions.
-        """
-        await self._task_controllers.pop(task_id).destroy()
+        function_executor_controller.add_task(
+            task=task_allocation.task,
+            allocation_id=task_allocation.allocation_id,
+        )
 
     def _valid_task_allocations(self, task_allocations: Iterable[TaskAllocation]):
         valid_task_allocations: List[TaskAllocation] = []
@@ -403,6 +400,7 @@ class ExecutorStateReconciler:
             validator = MessageValidator(task_allocation)
             try:
                 validator.required_field("function_executor_id")
+                validator.required_field("allocation_id")
             except ValueError as e:
                 # There's no way to report this error to Server so just log it.
                 logger.error(