indexify 0.4.26__tar.gz → 0.4.27__tar.gz

{indexify-0.4.26 → indexify-0.4.27}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.3
 Name: indexify
-Version: 0.4.26
+Version: 0.4.27
 Summary: Open Source Indexify components and helper tools
 Home-page: https://github.com/tensorlakeai/indexify
 License: Apache 2.0

{indexify-0.4.26 → indexify-0.4.27}/pyproject.toml

@@ -1,7 +1,7 @@
 [tool.poetry]
 name = "indexify"
 # Incremented if any of the components provided in this packages are updated.
-version = "0.4.26"
+version = "0.4.27"
 description = "Open Source Indexify components and helper tools"
 authors = ["Tensorlake Inc. <support@tensorlake.ai>"]
 license = "Apache 2.0"

{indexify-0.4.26 → indexify-0.4.27}/src/indexify/executor/function_executor/function_executor.py

@@ -308,10 +308,7 @@ async def _initialize_server(
     customer_code_timeout_sec: float,
     logger: Any,
 ) -> FunctionExecutorInitializationResult:
-    with (
-        metric_initialize_rpc_errors.count_exceptions(),
-        metric_initialize_rpc_latency.time(),
-    ):
+    with metric_initialize_rpc_latency.time():
         try:
             initialize_response: InitializeResponse = await stub.initialize(
                 initialize_request,
@@ -322,6 +319,9 @@ async def _initialize_server(
                 response=initialize_response,
             )
         except grpc.aio.AioRpcError as e:
+            # Increment the metric manually as we're not raising this exception.
+            metric_initialize_rpc_errors.inc()
+            metric_create_errors.inc()
             if e.code() == grpc.StatusCode.DEADLINE_EXCEEDED:
                 return FunctionExecutorInitializationResult(
                     is_timeout=True,

indexify-0.4.27/src/indexify/executor/function_executor_controller/aio_utils.py

@@ -0,0 +1,26 @@
+import asyncio
+from typing import Any, Optional
+
+
+async def shielded_await(task: asyncio.Task, logger: Any) -> Any:
+    """Awaits the supplied task and ignores cancellations until it's done.
+
+    Cancels itself if the task is cancelled once the task is done.
+    """
+    cancelled_error: Optional[asyncio.CancelledError] = None
+
+    while not task.done():
+        try:
+            # Shield to make sure that task is not cancelled.
+            await asyncio.shield(task)
+        except asyncio.CancelledError as e:
+            logger.info(
+                "ignoring aio task cancellation until it's finished",
+                task_name=task.get_name(),
+            )
+            cancelled_error = e
+
+    if cancelled_error is not None:
+        raise cancelled_error
+    else:
+        return task.result()

{indexify-0.4.26 → indexify-0.4.27}/src/indexify/executor/function_executor_controller/create_function_executor.py

@@ -25,6 +25,7 @@ from indexify.proto.executor_api_pb2 import (
     FunctionExecutorTerminationReason,
 )
 
+from .aio_utils import shielded_await
 from .downloads import download_graph
 from .events import FunctionExecutorCreated
 
@@ -82,8 +83,13 @@ async def create_function_executor(
         logger=logger,
     )
     if fe_created_event.function_executor is None:
+        # _to_fe_created_event doesn't like the FE, destroy it.
+        fe_destroy_task: asyncio.Task = asyncio.create_task(
+            function_executor.destroy(),
+            name=f"destroy function executor {function_executor_description.id}",
+        )
         try:
-            await asyncio.shield(function_executor.destroy())
+            await shielded_await(fe_destroy_task, logger)
         except asyncio.CancelledError:
             # destroy() finished due to the shield, return fe_created_event.
             pass
@@ -230,9 +236,16 @@ async def _create_function_executor(
             )
         )
         return (function_executor, result)
-    except BaseException:  # includes asyncio.CancelledError and anything else
+    except BaseException:
+        fe_destroy_task: asyncio.Task = asyncio.create_task(
+            function_executor.destroy(),
+            name=f"destroy function executor {function_executor_description.id}",
+        )
         # This await is a cancellation point, need to shield to ensure we destroyed the FE.
-        await asyncio.shield(function_executor.destroy())
+        await shielded_await(
+            fe_destroy_task,
+            logger,
+        )
         raise
 
 

{indexify-0.4.26 → indexify-0.4.27}/src/indexify/executor/function_executor_controller/function_executor_controller.py

@@ -407,7 +407,7 @@ class FunctionExecutorController:
         """Spawns an aio task for the supplied coroutine.
 
         The coroutine should return an event that will be added to the FE controller events.
-        The coroutine should not raise any exceptions including BaseException and asyncio.CancelledError.
+        The coroutine should not raise any exceptions including BaseException.
         on_exception event will be added to the FE controller events if the aio task raises an unexpected exception.
         on_exception is required to not silently stall the task processing due to an unexpected exception.
         If task_info is not None, the aio task will be associated with the task_info while the aio task is running.
@@ -422,6 +422,14 @@ class FunctionExecutorController:
         async def coroutine_wrapper() -> None:
             try:
                 self._add_event(await aio, source=aio_task_name)
+            except asyncio.CancelledError:
+                # Workaround for scenario when coroutine_wrapper gets cancelled at `await aio` before aio starts.
+                # In this case aio doesn't handle the cancellation and doesn't return the right event.
+                # A fix for this is to cancel aio instead of coroutine_wrapper. We'll need to keep
+                # references to both coroutine_wrapper and aio, cause event loop uses weak references to
+                # tasks. Not doing this for now. Using on_exception is good enough because not started aios don't
+                # need to do anything special on cancellation.
+                self._add_event(on_exception, source=aio_task_name)
             except BaseException as e:
                 logger.error(
                     "unexpected exception in aio task",

{indexify-0.4.26 → indexify-0.4.27}/src/indexify/executor/function_executor_controller/terminate_function_executor.py

@@ -4,6 +4,7 @@ from typing import Any, List, Optional
 from indexify.executor.function_executor.function_executor import FunctionExecutor
 from indexify.proto.executor_api_pb2 import FunctionExecutorTerminationReason
 
+from .aio_utils import shielded_await
 from .events import FunctionExecutorTerminated
 
 
@@ -18,7 +19,7 @@ async def terminate_function_executor(
 
     The supplied lock is used to ensure that if a destroy operation is in progress,
     then another caller won't return immediately assuming that the destroy is complete
-    due to its idempotency.
+    due to its idempotency. Ignores cancellations while destroying the function executor.
 
     Doesn't raise any exceptions.
     """
@@ -29,9 +30,12 @@ async def terminate_function_executor(
             logger.info(
                 "destroying function executor",
             )
+            fe_destroy_task: asyncio.Task = asyncio.create_task(
+                function_executor.destroy(),
+                name="destroy function executor",
+            )
             try:
-                # This await is a cancellation point, need to shield to ensure we destroyed the FE.
-                await asyncio.shield(function_executor.destroy())
+                await shielded_await(fe_destroy_task, logger)
             except asyncio.CancelledError:
                 # We actually destroyed the FE so we can return without error.
                 pass

{indexify-0.4.26 → indexify-0.4.27}/src/indexify/executor/state_reconciler.py

@@ -193,6 +193,8 @@ class ExecutorStateReconciler:
                     timeout=_DESIRED_EXECUTOR_STATES_TIMEOUT_SEC,
                 )
             except asyncio.TimeoutError:
+                # These log lines really help to debug networking issues. When there are
+                # no networking issues and the fleet is not idle we don't get excessive logging here.
                 self._logger.info(
                     f"No desired state received from Server within {_DESIRED_EXECUTOR_STATES_TIMEOUT_SEC} sec, recreating the stream to ensure it is healthy"
                 )