indexify 0.4.21__tar.gz → 0.4.22__tar.gz

{indexify-0.4.21 → indexify-0.4.22}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.3
 Name: indexify
-Version: 0.4.21
+Version: 0.4.22
 Summary: Open Source Indexify components and helper tools
 Home-page: https://github.com/tensorlakeai/indexify
 License: Apache 2.0
@@ -14,7 +14,7 @@ Classifier: Programming Language :: Python :: 3.11
 Classifier: Programming Language :: Python :: 3.12
 Classifier: Programming Language :: Python :: 3.13
 Requires-Dist: aiohttp (>=3.12.14,<4.0.0)
-Requires-Dist: boto3 (>=1.39.14,<2.0.0)
+Requires-Dist: boto3 (>=1.39.15,<2.0.0)
 Requires-Dist: docker (>=7.1.0,<8.0.0)
 Requires-Dist: httpx[http2] (==0.27.2)
 Requires-Dist: nanoid (>=2.0.0,<3.0.0)
@@ -22,7 +22,7 @@ Requires-Dist: prometheus-client (>=0.22.1,<0.23.0)
 Requires-Dist: psutil (>=7.0.0,<8.0.0)
 Requires-Dist: pydantic (>=2.11,<3.0)
 Requires-Dist: requests (>=2.32.4,<3.0.0)
-Requires-Dist: tensorlake (==0.2.33)
+Requires-Dist: tensorlake (==0.2.37)
 Requires-Dist: urllib3 (>=2.5.0,<3.0.0)
 Project-URL: Repository, https://github.com/tensorlakeai/indexify
 Description-Content-Type: text/markdown

{indexify-0.4.21 → indexify-0.4.22}/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.21"
+version = "0.4.22"
 description = "Open Source Indexify components and helper tools"
 authors = ["Tensorlake Inc. <support@tensorlake.ai>"]
 license = "Apache 2.0"
@@ -26,10 +26,10 @@ httpx = { version = "0.27.2", extras = ["http2"] }
 pydantic = "^2.11"
 prometheus-client = "^0.22.1"
 psutil = "^7.0.0"
-boto3 = "^1.39.14"
+boto3 = "^1.39.15"
 # Adds function-executor binary, utils lib, sdk used in indexify-cli commands.
 # We need to specify the tensorlake version exactly because pip install doesn't respect poetry.lock files.
-tensorlake = "0.2.33"
+tensorlake = "0.2.37"
 # Uncomment the next line to use local tensorlake package (only for development!)
 # tensorlake = { path = "../tensorlake", develop = true }
 # grpcio is provided by tensorlake

{indexify-0.4.21 → indexify-0.4.22}/src/indexify/executor/function_executor_controller/metrics/run_task.py

@@ -6,23 +6,24 @@ from indexify.executor.monitoring.metrics import (
 
 metric_function_executor_run_task_rpcs: prometheus_client.Counter = (
     prometheus_client.Counter(
-        "function_executor_run_task_rpcs", "Number of Function Executor run task RPCs"
+        "function_executor_run_task_rpcs",
+        "Number of Function Executor run task lifecycle RPC sequences",
     )
 )
 metric_function_executor_run_task_rpc_errors: prometheus_client.Counter = (
     prometheus_client.Counter(
         "function_executor_run_task_rpc_errors",
-        "Number of Function Executor run task RPC errors",
+        "Number of Function Executor run task lifecycle RPC errors",
     )
 )
 metric_function_executor_run_task_rpc_latency: prometheus_client.Histogram = (
     latency_metric_for_customer_controlled_operation(
-        "function_executor_run_task_rpc", "Function Executor run task RPC"
+        "function_executor_run_task_rpc", "Function Executor run task lifecycle RPC"
     )
 )
 metric_function_executor_run_task_rpcs_in_progress: prometheus_client.Gauge = (
     prometheus_client.Gauge(
         "function_executor_run_task_rpcs_in_progress",
-        "Number of Function Executor run task RPCs in progress",
+        "Number of Function Executor run task lifecycle RPCs in progress",
     )
 )

{indexify-0.4.21 → indexify-0.4.22}/src/indexify/executor/function_executor_controller/run_task.py

@@ -6,9 +6,13 @@ from typing import Any, Optional
 
 import grpc
 from tensorlake.function_executor.proto.function_executor_pb2 import (
-    RunTaskRequest,
-    RunTaskResponse,
+    AwaitTaskProgress,
+    AwaitTaskRequest,
+    CreateTaskRequest,
+    DeleteTaskRequest,
+    FunctionInputs,
     SerializedObject,
+    Task,
 )
 from tensorlake.function_executor.proto.function_executor_pb2 import (
     TaskFailureReason as FETaskFailureReason,
@@ -16,6 +20,9 @@ from tensorlake.function_executor.proto.function_executor_pb2 import (
 from tensorlake.function_executor.proto.function_executor_pb2 import (
     TaskOutcomeCode as FETaskOutcomeCode,
 )
+from tensorlake.function_executor.proto.function_executor_pb2 import (
+    TaskResult,
+)
 from tensorlake.function_executor.proto.function_executor_pb2_grpc import (
     FunctionExecutorStub,
 )
@@ -44,6 +51,9 @@ _ENABLE_INJECT_TASK_CANCELLATIONS = (
     os.getenv("INDEXIFY_INJECT_TASK_CANCELLATIONS", "0") == "1"
 )
 
+_CREATE_TASK_TIMEOUT_SECS = 5
+_DELETE_TASK_TIMEOUT_SECS = 5
+
 
 async def run_task_on_function_executor(
     task_info: TaskInfo, function_executor: FunctionExecutor, logger: Any
@@ -53,21 +63,21 @@ async def run_task_on_function_executor(
     Doesn't raise any exceptions.
     """
     logger = logger.bind(module=__name__)
-    request: RunTaskRequest = RunTaskRequest(
+    task = Task(
+        task_id=task_info.allocation.task.id,
         namespace=task_info.allocation.task.namespace,
         graph_name=task_info.allocation.task.graph_name,
         graph_version=task_info.allocation.task.graph_version,
         function_name=task_info.allocation.task.function_name,
         graph_invocation_id=task_info.allocation.task.graph_invocation_id,
-        task_id=task_info.allocation.task.id,
         allocation_id=task_info.allocation.allocation_id,
-        function_input=task_info.input,
+        request=FunctionInputs(function_input=task_info.input),
     )
     # Don't keep the input in memory after we started running the task.
     task_info.input = None
 
     if task_info.init_value is not None:
-        request.function_init_value.CopyFrom(task_info.init_value)
+        task.request.function_init_value.CopyFrom(task_info.init_value)
         # Don't keep the init value in memory after we started running the task.
         task_info.init_value = None
 
@@ -78,50 +88,75 @@ async def run_task_on_function_executor(
 
     metric_function_executor_run_task_rpcs.inc()
     metric_function_executor_run_task_rpcs_in_progress.inc()
-    start_time = time.monotonic()
     # Not None if the Function Executor should be terminated after running the task.
     function_executor_termination_reason: Optional[
         FunctionExecutorTerminationReason
     ] = None
-    execution_start_time: Optional[float] = None
+
+    # NB: We start this timer before invoking the first RPC, since
+    # user code should be executing by the time the create_task() RPC
+    # returns, so not attributing the task management RPC overhead to
+    # the user would open a possibility for abuse. (This is somewhat
+    # mitigated by the fact that these RPCs should have a very low
+    # overhead.)
+    execution_start_time: Optional[float] = time.monotonic()
 
     # If this RPC failed due to customer code crashing the server we won't be
     # able to detect this. We'll treat this as our own error for now and thus
     # let the AioRpcError to be raised here.
     timeout_sec = task_info.allocation.task.timeout_ms / 1000.0
     try:
-        channel: grpc.aio.Channel = function_executor.channel()
-        execution_start_time = time.monotonic()
-        response: RunTaskResponse = await FunctionExecutorStub(channel).run_task(
-            request, timeout=timeout_sec
-        )
-        task_info.output = _task_output_from_function_executor_response(
+        task_result = await _run_task_rpcs(task, function_executor, timeout_sec)
+
+        task_info.output = _task_output_from_function_executor_result(
             allocation=task_info.allocation,
-            response=response,
+            result=task_result,
             execution_start_time=execution_start_time,
             execution_end_time=time.monotonic(),
             logger=logger,
         )
+    except asyncio.TimeoutError:
+        # This is an await_task() RPC timeout - we're not getting
+        # progress messages or a task completion.
+        function_executor_termination_reason = (
+            FunctionExecutorTerminationReason.FUNCTION_EXECUTOR_TERMINATION_REASON_FUNCTION_TIMEOUT
+        )
+        task_info.output = TaskOutput.function_timeout(
+            allocation=task_info.allocation,
+            timeout_sec=timeout_sec,
+            execution_start_time=execution_start_time,
+            execution_end_time=time.monotonic(),
+        )
     except grpc.aio.AioRpcError as e:
+        # This indicates some sort of problem communicating with the FE.
+        #
+        # NB: We charge the user in these situations: code within the
+        # FE is not isolated, so not charging would enable abuse.
+        #
+        # This is an unexpected situation, though, so we make sure to
+        # log the situation for further investigation.
+
+        function_executor_termination_reason = (
+            FunctionExecutorTerminationReason.FUNCTION_EXECUTOR_TERMINATION_REASON_UNHEALTHY
+        )
+        metric_function_executor_run_task_rpc_errors.inc()
+
         if e.code() == grpc.StatusCode.DEADLINE_EXCEEDED:
-            # The task is still running in FE, we only cancelled the client-side RPC.
-            function_executor_termination_reason = (
-                FunctionExecutorTerminationReason.FUNCTION_EXECUTOR_TERMINATION_REASON_FUNCTION_TIMEOUT
-            )
-            task_info.output = TaskOutput.function_timeout(
-                allocation=task_info.allocation,
-                timeout_sec=timeout_sec,
-                execution_start_time=execution_start_time,
-                execution_end_time=time.monotonic(),
-            )
+            # This is either a create_task() RPC timeout or a
+            # delete_task() RPC timeout; either suggests that the FE
+            # is unhealthy.
+            logger.error("task management RPC execution deadline exceeded", exc_info=e)
         else:
-            metric_function_executor_run_task_rpc_errors.inc()
-            logger.error("task execution failed", exc_info=e)
-            task_info.output = TaskOutput.internal_error(
-                allocation=task_info.allocation,
-                execution_start_time=execution_start_time,
-                execution_end_time=time.monotonic(),
-            )
+            # This is a status from an unsuccessful RPC; this
+            # shouldn't happen, but we handle it.
+            logger.error("task management RPC failed", exc_info=e)
+
+        task_info.output = TaskOutput.function_executor_unresponsive(
+            allocation=task_info.allocation,
+            execution_start_time=execution_start_time,
+            execution_end_time=time.monotonic(),
+        )
+
     except asyncio.CancelledError:
         # The task is still running in FE, we only cancelled the client-side RPC.
         function_executor_termination_reason = (
@@ -133,15 +168,20 @@ async def run_task_on_function_executor(
             execution_end_time=time.monotonic(),
         )
     except Exception as e:
-        metric_function_executor_run_task_rpc_errors.inc()
-        logger.error("task execution failed", exc_info=e)
+        # This is an unexpected exception; we believe that this
+        # indicates an internal error.
+        logger.error(
+            "Unexpected internal error during task lifecycle RPC sequence", exc_info=e
+        )
         task_info.output = TaskOutput.internal_error(
             allocation=task_info.allocation,
             execution_start_time=execution_start_time,
             execution_end_time=time.monotonic(),
         )
 
-    metric_function_executor_run_task_rpc_latency.observe(time.monotonic() - start_time)
+    metric_function_executor_run_task_rpc_latency.observe(
+        time.monotonic() - execution_start_time
+    )
     metric_function_executor_run_task_rpcs_in_progress.dec()
 
     function_executor.invocation_state_client().remove_task_to_invocation_id_entry(
@@ -171,26 +211,78 @@ async def run_task_on_function_executor(
     )
 
 
-def _task_output_from_function_executor_response(
+async def _run_task_rpcs(
+    task: Task, function_executor: FunctionExecutor, timeout_sec: float
+) -> TaskResult:
+    """Runs the task, returning the result, reporting errors via exceptions."""
+
+    response: AwaitTaskProgress
+    channel: grpc.aio.Channel = function_executor.channel()
+    fe_stub = FunctionExecutorStub(channel)
+
+    # Create task with timeout
+    await fe_stub.create_task(
+        CreateTaskRequest(task=task), timeout=_CREATE_TASK_TIMEOUT_SECS
+    )
+
+    # Await task with timeout resets on each response
+    await_rpc = fe_stub.await_task(AwaitTaskRequest(task_id=task.task_id))
+
+    try:
+        while True:
+            # Wait for next response with fresh timeout each time
+            response = await asyncio.wait_for(await_rpc.read(), timeout=timeout_sec)
+            if response.WhichOneof("response") == "task_result":
+                # We're done waiting.
+                break
+
+            # NB: We don't actually check for other message types
+            # here; any message from the FE is treated as an
+            # indication that it's making forward progress.
+
+            if response == grpc.aio.EOF:
+                # Protocol error: we should get a task_result before
+                # we see the RPC complete.
+                raise grpc.aio.AioRpcError(
+                    grpc.StatusCode.CANCELLED,
+                    None,
+                    None,
+                    "Function Executor didn't return function/task alloc response",
+                )
+    finally:
+        # Cancel the outstanding RPC to ensure any resources in use
+        # are cleaned up; note that this is idempotent (in case the
+        # RPC has already completed).
+        await_rpc.cancel()
+
+    # Delete task with timeout
+    await fe_stub.delete_task(
+        DeleteTaskRequest(task_id=task.task_id), timeout=_DELETE_TASK_TIMEOUT_SECS
+    )
+
+    return response.task_result
+
+
+def _task_output_from_function_executor_result(
     allocation: TaskAllocation,
-    response: RunTaskResponse,
+    result: TaskResult,
     execution_start_time: Optional[float],
     execution_end_time: Optional[float],
     logger: Any,
 ) -> TaskOutput:
-    response_validator = MessageValidator(response)
+    response_validator = MessageValidator(result)
     response_validator.required_field("stdout")
     response_validator.required_field("stderr")
     response_validator.required_field("outcome_code")
 
     metrics = TaskMetrics(counters={}, timers={})
-    if response.HasField("metrics"):
+    if result.HasField("metrics"):
         # Can be None if e.g. function failed.
-        metrics.counters = dict(response.metrics.counters)
-        metrics.timers = dict(response.metrics.timers)
+        metrics.counters = dict(result.metrics.counters)
+        metrics.timers = dict(result.metrics.timers)
 
     outcome_code: TaskOutcomeCode = _to_task_outcome_code(
-        response.outcome_code, logger=logger
+        result.outcome_code, logger=logger
     )
     failure_reason: Optional[TaskFailureReason] = None
     invocation_error_output: Optional[SerializedObject] = None
@@ -198,11 +290,11 @@ def _task_output_from_function_executor_response(
     if outcome_code == TaskOutcomeCode.TASK_OUTCOME_CODE_FAILURE:
         response_validator.required_field("failure_reason")
         failure_reason: Optional[TaskFailureReason] = _to_task_failure_reason(
-            response.failure_reason, logger
+            result.failure_reason, logger
         )
         if failure_reason == TaskFailureReason.TASK_FAILURE_REASON_INVOCATION_ERROR:
             response_validator.required_field("invocation_error_output")
-            invocation_error_output = response.invocation_error_output
+            invocation_error_output = result.invocation_error_output
 
     if _ENABLE_INJECT_TASK_CANCELLATIONS:
         logger.warning("injecting cancellation failure for the task allocation")
@@ -217,10 +309,10 @@ def _task_output_from_function_executor_response(
         outcome_code=outcome_code,
         failure_reason=failure_reason,
         invocation_error_output=invocation_error_output,
-        function_outputs=response.function_outputs,
-        next_functions=response.next_functions,
-        stdout=response.stdout,
-        stderr=response.stderr,
+        function_outputs=result.function_outputs,
+        next_functions=result.next_functions,
+        stdout=result.stdout,
+        stderr=result.stderr,
         metrics=metrics,
         execution_start_time=execution_start_time,
         execution_end_time=execution_end_time,

{indexify-0.4.21 → indexify-0.4.22}/src/indexify/executor/function_executor_controller/task_output.py

@@ -95,6 +95,23 @@ class TaskOutput:
             execution_end_time=execution_end_time,
         )
 
+    @classmethod
+    def function_executor_unresponsive(
+        cls,
+        allocation: TaskAllocation,
+        execution_start_time: Optional[float],
+        execution_end_time: Optional[float],
+    ) -> "TaskOutput":
+        """Creates a TaskOutput for an unresponsive FE."""
+        # Task stdout, stderr is not available.
+        return TaskOutput(
+            allocation=allocation,
+            outcome_code=TaskOutcomeCode.TASK_OUTCOME_CODE_FAILURE,
+            failure_reason=TaskFailureReason.TASK_FAILURE_REASON_FUNCTION_ERROR,
+            execution_start_time=execution_start_time,
+            execution_end_time=execution_end_time,
+        )
+
     @classmethod
     def task_cancelled(
         cls,