indexify 0.4.21__py3-none-any.whl → 0.4.22__py3-none-any.whl

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/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/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,

{indexify-0.4.21.dist-info → indexify-0.4.22.dist-info}/METADATA

@@ -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.dist-info → indexify-0.4.22.dist-info}/RECORD

@@ -34,12 +34,12 @@ indexify/executor/function_executor_controller/message_validators.py,sha256=aNiZ
 indexify/executor/function_executor_controller/metrics/completed_task_metrics.py,sha256=53EGBCLwCEV-RBBeyLPTElrtcveaEM0Fwxs9NmC1Hn8,2724
 indexify/executor/function_executor_controller/metrics/downloads.py,sha256=KOVTE2OZPCewnCooPyCK1maKe9ddMPvBFp7D_igqugQ,2708
 indexify/executor/function_executor_controller/metrics/function_executor_controller.py,sha256=gyIZHbdsNSnrA6z4WDaRAunNolFrbbg1pa8JaL_ODNE,2666
-indexify/executor/function_executor_controller/metrics/run_task.py,sha256=kdjw-Bi788B8Pq4eL08k8HVtcjIA3y8pLLSNWRBI8YA,1006
+indexify/executor/function_executor_controller/metrics/run_task.py,sha256=ZFv_nw5_pKUJoTaavSyzdglQKW4uvC2XyK8S6xi9xLQ,1064
 indexify/executor/function_executor_controller/metrics/upload_task_output.py,sha256=Ppf8NujCNbQzFelJiuh8Sutcjty7hnkFz1dWQLYIgQI,1464
 indexify/executor/function_executor_controller/prepare_task.py,sha256=AKbo_H_5pOKdxFKKkzdOb1WhQ0XT-4Qm9D3iIsukyMU,1247
-indexify/executor/function_executor_controller/run_task.py,sha256=IqJU7EDu5iacvlT8A17-_78k7Mt8bUjillxWjsr-_A4,11293
+indexify/executor/function_executor_controller/run_task.py,sha256=2qpP6YwKoi_sEWDJIJ43EyBgO3Z3yYDVFpncB01DQf4,14586
 indexify/executor/function_executor_controller/task_info.py,sha256=ZEdypd8QVmYbrLt1186Ed9YEQwrO0Sx_hKH0QLg1DVY,1181
-indexify/executor/function_executor_controller/task_output.py,sha256=rSUghdPTfWsYQ7h1ZgmPOp7VTPAD4fPWkpN7MXsmlV0,6702
+indexify/executor/function_executor_controller/task_output.py,sha256=krwJC8GgVWYM_b_v5dHtVytsH3mHz3Qi7REiEM2D7nA,7358
 indexify/executor/function_executor_controller/terminate_function_executor.py,sha256=YLDlKoanfUBcy7A9ydCYdUsDwApjcTTn1o4tjNVN_QA,1281
 indexify/executor/function_executor_controller/upload_task_output.py,sha256=fEZm5eodx5rNLQYFhmdkMDD9qjX3_wKo64x4aUKTu34,10403
 indexify/executor/host_resources/host_resources.py,sha256=eUyP05EX7QdOtQ5vbX_KCpvnBS2B7fl06UWeF9Oigns,3813
@@ -64,7 +64,7 @@ indexify/proto/executor_api.proto,sha256=Zvx5eTz_QFHCX2udzOwHZD-ncJfDlXeKVbNBhsa
 indexify/proto/executor_api_pb2.py,sha256=iSiO2Aw7vQV5Jp7IoSeIC8q32WjE-s-AZ0Hc5MwI0uI,16266
 indexify/proto/executor_api_pb2.pyi,sha256=oReajQdT2yNPS4kZd3gDMBdbseHv-gU15fpeT0v3ZBs,22819
 indexify/proto/executor_api_pb2_grpc.py,sha256=u9GEQV4nm_GvApRxjVo806CkgBMBVReb5IVrcaDaliY,7520
-indexify-0.4.21.dist-info/METADATA,sha256=UpSK8Rpb61jT_7IqUUBjL4NXWjk5f7jELuFo-YZLzJE,1354
-indexify-0.4.21.dist-info/WHEEL,sha256=RaoafKOydTQ7I_I3JTrPCg6kUmTgtm4BornzOqyEfJ8,88
-indexify-0.4.21.dist-info/entry_points.txt,sha256=rMJqbE5KPZIXTPIfAtVIM4zpUElqYVgEYd6i7N23zzg,49
-indexify-0.4.21.dist-info/RECORD,,
+indexify-0.4.22.dist-info/METADATA,sha256=poiBPNliE54kWJyR5_FSISz-gWYYn6wXIzmgmn-Es3k,1354
+indexify-0.4.22.dist-info/WHEEL,sha256=RaoafKOydTQ7I_I3JTrPCg6kUmTgtm4BornzOqyEfJ8,88
+indexify-0.4.22.dist-info/entry_points.txt,sha256=rMJqbE5KPZIXTPIfAtVIM4zpUElqYVgEYd6i7N23zzg,49
+indexify-0.4.22.dist-info/RECORD,,