indexify 0.4.28__py3-none-any.whl → 0.4.30__py3-none-any.whl

indexify/executor/function_executor_controller/loggers.py

@@ -49,7 +49,21 @@ def task_allocation_logger(task_allocation: TaskAllocation, logger: Any) -> Any:
     Doesn't assume that the supplied TaskAllocation is valid.
     """
     if task_allocation.HasField("task"):
-        logger = _task_logger(task_allocation.task, logger)
+        task: Task = task_allocation.task
+        logger = logger.bind(
+            task_id=task.id if task.HasField("id") else None,
+            namespace=task.namespace if task.HasField("namespace") else None,
+            graph=task.graph_name if task.HasField("graph_name") else None,
+            graph_version=(
+                task.graph_version if task.HasField("graph_version") else None
+            ),
+            fn=task.function_name if task.HasField("function_name") else None,
+            invocation_id=(
+                task.graph_invocation_id
+                if task.HasField("graph_invocation_id")
+                else None
+            ),
+        )
     return logger.bind(
         allocation_id=(
             task_allocation.allocation_id
@@ -87,19 +101,3 @@ def task_result_logger(task_result: TaskResult, logger: Any) -> Any:
             else None
         ),
     )
-
-
-def _task_logger(task: Task, logger: Any) -> Any:
-    """Returns a logger bound with the task's metadata.
-
-    The function assumes that the task might be invalid."""
-    return logger.bind(
-        task_id=task.id if task.HasField("id") else None,
-        namespace=task.namespace if task.HasField("namespace") else None,
-        graph=task.graph_name if task.HasField("graph_name") else None,
-        graph_version=task.graph_version if task.HasField("graph_version") else None,
-        fn=task.function_name if task.HasField("function_name") else None,
-        invocation_id=(
-            task.graph_invocation_id if task.HasField("graph_invocation_id") else None
-        ),
-    )

indexify/executor/function_executor_controller/message_validators.py

@@ -3,7 +3,6 @@ from tensorlake.function_executor.proto.message_validator import MessageValidato
 from indexify.proto.executor_api_pb2 import (
     DataPayload,
     FunctionExecutorDescription,
-    Task,
     TaskAllocation,
 )
 
@@ -50,15 +49,8 @@ def validate_task_allocation(task_allocation: TaskAllocation) -> None:
     validator.required_field("allocation_id")
     if not task_allocation.HasField("task"):
         raise ValueError("TaskAllocation must have a 'task' field.")
-    _validate_task(task_allocation.task)
 
-
-def _validate_task(task: Task) -> None:
-    """Validates the supplied Task.
-
-    Raises ValueError if the Task is not valid.
-    """
-    validator = MessageValidator(task)
+    validator = MessageValidator(task_allocation.task)
     validator.required_field("id")
     validator.required_field("namespace")
     validator.required_field("graph_name")
@@ -70,9 +62,9 @@ def _validate_task(task: Task) -> None:
     validator.required_field("output_payload_uri_prefix")
     validator.required_field("retry_policy")
 
-    _validate_data_payload(task.input)
-    if task.HasField("reducer_input"):
-        _validate_data_payload(task.reducer_input)
+    _validate_data_payload(task_allocation.task.input)
+    if task_allocation.task.HasField("reducer_input"):
+        _validate_data_payload(task_allocation.task.reducer_input)
 
 
 def _validate_data_payload(data_payload: DataPayload) -> None:

indexify/executor/function_executor_controller/metrics/completed_task_allocation_metrics.py

@@ -0,0 +1,70 @@
+import prometheus_client
+
+from indexify.executor.monitoring.metrics import (
+    latency_metric_for_customer_controlled_operation,
+)
+
+metric_task_allocations_completed: prometheus_client.Counter = (
+    prometheus_client.Counter(
+        "task_allocations_completed",
+        "Number of task allocations that were completed",
+        ["outcome_code", "failure_reason"],
+    )
+)
+METRIC_TASK_ALLOCATIONS_COMPLETED_OUTCOME_CODE_ALL = "all"
+METRIC_TASK_ALLOCATIONS_COMPLETED_OUTCOME_CODE_SUCCESS = "success"
+METRIC_TASK_ALLOCATIONS_COMPLETED_OUTCOME_CODE_FAILURE = "failure"
+
+METRIC_TASK_ALLOCATIONS_COMPLETED_FAILURE_REASON_ALL = "all"
+# Used when the task is successfull.
+METRIC_TASK_ALLOCATIONS_COMPLETED_FAILURE_REASON_NONE = "none"
+# Matches TASK_FAILURE_REASON_UNKNOWN
+METRIC_TASK_ALLOCATIONS_COMPLETED_FAILURE_REASON_UNKNOWN = "unknown"
+# Includes all function errors including timeouts to reduce cardinality.
+METRIC_TASK_ALLOCATIONS_COMPLETED_FAILURE_REASON_FUNCTION_ERROR = "function_error"
+# Includes all internal errors to reduce cardinality.
+METRIC_TASK_ALLOCATIONS_COMPLETED_FAILURE_REASON_INTERNAL_ERROR = "internal_error"
+# Matches TASK_FAILURE_REASON_TASK_CANCELLED
+METRIC_TASK_ALLOCATIONS_COMPLETED_FAILURE_REASON_TASK_CANCELLED = "task_cancelled"
+# Matches TASK_FAILURE_REASON_FUNCTION_EXECUTOR_TERMINATED
+METRIC_TASK_ALLOCATIONS_COMPLETED_FAILURE_REASON_FUNCTION_EXECUTOR_TERMINATED = (
+    "function_executor_terminated"
+)
+
+# Valid combinations of the labels:
+metric_task_allocations_completed.labels(
+    outcome_code=METRIC_TASK_ALLOCATIONS_COMPLETED_OUTCOME_CODE_ALL,
+    failure_reason=METRIC_TASK_ALLOCATIONS_COMPLETED_FAILURE_REASON_ALL,
+)
+metric_task_allocations_completed.labels(
+    outcome_code=METRIC_TASK_ALLOCATIONS_COMPLETED_OUTCOME_CODE_SUCCESS,
+    failure_reason=METRIC_TASK_ALLOCATIONS_COMPLETED_FAILURE_REASON_NONE,
+)
+
+metric_task_allocations_completed.labels(
+    outcome_code=METRIC_TASK_ALLOCATIONS_COMPLETED_OUTCOME_CODE_FAILURE,
+    failure_reason=METRIC_TASK_ALLOCATIONS_COMPLETED_FAILURE_REASON_UNKNOWN,
+)
+metric_task_allocations_completed.labels(
+    outcome_code=METRIC_TASK_ALLOCATIONS_COMPLETED_OUTCOME_CODE_FAILURE,
+    failure_reason=METRIC_TASK_ALLOCATIONS_COMPLETED_FAILURE_REASON_FUNCTION_ERROR,
+)
+metric_task_allocations_completed.labels(
+    outcome_code=METRIC_TASK_ALLOCATIONS_COMPLETED_OUTCOME_CODE_FAILURE,
+    failure_reason=METRIC_TASK_ALLOCATIONS_COMPLETED_FAILURE_REASON_INTERNAL_ERROR,
+)
+metric_task_allocations_completed.labels(
+    outcome_code=METRIC_TASK_ALLOCATIONS_COMPLETED_OUTCOME_CODE_FAILURE,
+    failure_reason=METRIC_TASK_ALLOCATIONS_COMPLETED_FAILURE_REASON_TASK_CANCELLED,
+)
+metric_task_allocations_completed.labels(
+    outcome_code=METRIC_TASK_ALLOCATIONS_COMPLETED_OUTCOME_CODE_FAILURE,
+    failure_reason=METRIC_TASK_ALLOCATIONS_COMPLETED_FAILURE_REASON_FUNCTION_EXECUTOR_TERMINATED,
+)
+
+metric_task_allocation_completion_latency: prometheus_client.Histogram = (
+    latency_metric_for_customer_controlled_operation(
+        "task_allocation_completion",
+        "task allocation completion from the moment it got fetched until its output got uploaded to blob store",
+    )
+)

indexify/executor/function_executor_controller/metrics/finalize_task_allocation.py

@@ -0,0 +1,26 @@
+import prometheus_client
+
+from indexify.executor.monitoring.metrics import latency_metric_for_fast_operation
+
+# Task allocation finalization metrics.
+metric_task_allocation_finalizations: prometheus_client.Counter = (
+    prometheus_client.Counter(
+        "task_allocation_finalizations",
+        "Number of task allocation finalizations",
+    )
+)
+metric_task_allocation_finalization_errors: prometheus_client.Counter = (
+    prometheus_client.Counter(
+        "task_allocation_finalization_errors",
+        "Number of task allocation finalization errors",
+    )
+)
+metric_task_allocations_finalizing: prometheus_client.Gauge = prometheus_client.Gauge(
+    "task_allocations_finalizing",
+    "Number of task allocations currently finalizing",
+)
+metric_task_allocation_finalization_latency: prometheus_client.Histogram = (
+    latency_metric_for_fast_operation(
+        "task_allocation_finalization", "task allocation finalization"
+    )
+)

indexify/executor/function_executor_controller/metrics/function_executor_controller.py

@@ -12,24 +12,25 @@ metric_control_loop_handle_event_latency: prometheus_client.Histogram = (
     )
 )
 
-metric_tasks_fetched: prometheus_client.Counter = prometheus_client.Counter(
-    "tasks_fetched", "Number of tasks that were fetched from Server"
+metric_task_allocations_fetched: prometheus_client.Counter = prometheus_client.Counter(
+    "task_allocations_fetched",
+    "Number of task allocations that were fetched from Server",
 )
 
-metric_schedule_task_latency: prometheus_client.Histogram = (
+metric_schedule_task_allocation_latency: prometheus_client.Histogram = (
     latency_metric_for_customer_controlled_operation(
-        "schedule_task",
-        "Schedule a task for execution after it got ready for execution",
+        "schedule_task_allocation",
+        "Schedule a task allocation for execution after it got ready for execution",
     )
 )
-metric_runnable_tasks: prometheus_client.Gauge = prometheus_client.Gauge(
-    "runnable_tasks",
-    "Number of tasks that are ready for execution but are waiting to get scheduled to run on Function Executor (typically waiting for a free Function Executor)",
+metric_runnable_task_allocations: prometheus_client.Gauge = prometheus_client.Gauge(
+    "runnable_task_allocations",
+    "Number of task allocations that are ready for execution but are waiting to get scheduled to run on Function Executor (typically waiting for a free Function Executor)",
 )
-metric_runnable_tasks_per_function_name: prometheus_client.Gauge = (
+metric_runnable_task_allocations_per_function_name: prometheus_client.Gauge = (
     prometheus_client.Gauge(
-        "runnable_tasks_per_function_name",
-        "Number of tasks that are ready for execution but are waiting to get scheduled to run on Function Executor (typically waiting for a free Function Executor)",
+        "runnable_task_allocations_per_function_name",
+        "Number of task allocations that are ready for execution but are waiting to get scheduled to run on Function Executor (typically waiting for a free Function Executor)",
         ["function_name"],
     )
 )

indexify/executor/function_executor_controller/metrics/prepare_task_allocation.py

@@ -0,0 +1,27 @@
+import prometheus_client
+
+from indexify.executor.monitoring.metrics import latency_metric_for_fast_operation
+
+metric_task_allocation_preparations: prometheus_client.Counter = (
+    prometheus_client.Counter(
+        "task_allocation_preparations",
+        "Number of task allocation preparations for execution",
+    )
+)
+metric_task_allocation_preparation_errors: prometheus_client.Counter = (
+    prometheus_client.Counter(
+        "task_allocation_preparation_errors",
+        "Number of task allocation preparation errors",
+    )
+)
+metric_task_allocation_preparation_latency: prometheus_client.Histogram = (
+    latency_metric_for_fast_operation(
+        "task_allocation_preparation", "task allocation preparation for execution"
+    )
+)
+metric_task_allocations_getting_prepared: prometheus_client.Gauge = (
+    prometheus_client.Gauge(
+        "task_allocations_getting_prepared",
+        "Number of task allocations currently getting prepared for execution",
+    )
+)

indexify/executor/function_executor_controller/{prepare_task.py → prepare_task_allocation.py}

@@ -13,15 +13,15 @@ from indexify.executor.blob_store.blob_store import BLOBStore
 from indexify.proto.executor_api_pb2 import DataPayload, Task
 
 from .downloads import serialized_object_manifest_from_data_payload_proto
-from .events import TaskPreparationFinished
-from .metrics.prepare_task import (
-    metric_task_preparation_errors,
-    metric_task_preparation_latency,
-    metric_task_preparations,
-    metric_tasks_getting_prepared,
+from .events import TaskAllocationPreparationFinished
+from .metrics.prepare_task_allocation import (
+    metric_task_allocation_preparation_errors,
+    metric_task_allocation_preparation_latency,
+    metric_task_allocation_preparations,
+    metric_task_allocations_getting_prepared,
 )
-from .task_info import TaskInfo
-from .task_input import TaskInput
+from .task_allocation_info import TaskAllocationInfo
+from .task_allocation_input import TaskAllocationInput
 
 # The following constants are subject to S3 limits,
 # see https://docs.aws.amazon.com/AmazonS3/latest/userguide/qfacts.html.
@@ -42,9 +42,9 @@ _OUTPUT_BLOB_SLOWER_CHUNKS_COUNT: int = 100
 _INVOCATION_ERROR_MAX_SIZE_BYTES: int = 10 * 1024 * 1024  # 10 MB
 
 
-async def prepare_task(
-    task_info: TaskInfo, blob_store: BLOBStore, logger: Any
-) -> TaskPreparationFinished:
+async def prepare_task_allocation(
+    alloc_info: TaskAllocationInfo, blob_store: BLOBStore, logger: Any
+) -> TaskAllocationPreparationFinished:
     """Prepares the task for execution.
 
     If successful then the task is runnable.
@@ -54,43 +54,47 @@ async def prepare_task(
     start_time = time.monotonic()
     try:
         with (
-            metric_task_preparation_errors.count_exceptions(),
-            metric_tasks_getting_prepared.track_inprogress(),
-            metric_task_preparation_latency.time(),
+            metric_task_allocation_preparation_errors.count_exceptions(),
+            metric_task_allocations_getting_prepared.track_inprogress(),
+            metric_task_allocation_preparation_latency.time(),
         ):
-            metric_task_preparations.inc()
-            task_info.input = await _prepare_task_input(
-                task_info=task_info,
+            metric_task_allocation_preparations.inc()
+            alloc_info.input = await _prepare_task_alloc_input(
+                alloc_info=alloc_info,
                 blob_store=blob_store,
                 logger=logger,
             )
             logger.info(
-                "task was prepared for execution",
+                "task allocation was prepared for execution",
                 duration=time.monotonic() - start_time,
             )
-            return TaskPreparationFinished(
-                task_info=task_info,
+            return TaskAllocationPreparationFinished(
+                alloc_info=alloc_info,
                 is_success=True,
             )
     except asyncio.CancelledError:
-        return TaskPreparationFinished(task_info=task_info, is_success=False)
+        return TaskAllocationPreparationFinished(
+            alloc_info=alloc_info, is_success=False
+        )
     except BaseException as e:
         logger.error(
-            "failed to prepare task for execution",
+            "failed to prepare task allocation for execution",
             exc_info=e,
             duration=time.monotonic() - start_time,
         )
-        return TaskPreparationFinished(task_info=task_info, is_success=False)
+        return TaskAllocationPreparationFinished(
+            alloc_info=alloc_info, is_success=False
+        )
 
 
-async def _prepare_task_input(
-    task_info: TaskInfo, blob_store: BLOBStore, logger: Any
-) -> TaskInput:
+async def _prepare_task_alloc_input(
+    alloc_info: TaskAllocationInfo, blob_store: BLOBStore, logger: Any
+) -> TaskAllocationInput:
     """Prepares the task for execution.
 
     Raises an exception on error.
     """
-    task: Task = task_info.allocation.task
+    task: Task = alloc_info.allocation.task
     function_init_value_blob: Optional[BLOB] = None
     function_init_value: Optional[SerializedObjectInsideBLOB] = None
     if task.HasField("reducer_input"):
@@ -102,7 +106,7 @@ async def _prepare_task_input(
         function_init_value = _to_serialized_object_inside_blob(task.reducer_input)
 
     function_outputs_blob_uri: str = (
-        f"{task.output_payload_uri_prefix}.{task_info.allocation.allocation_id}.output"
+        f"{task.output_payload_uri_prefix}.{alloc_info.allocation.allocation_id}.output"
     )
     invocation_error_blob_uri: str = (
         f"{task.invocation_error_payload_uri_prefix}.{task.graph_invocation_id}.inverr"
@@ -136,7 +140,7 @@ async def _prepare_task_input(
             )
         raise
 
-    return TaskInput(
+    return TaskAllocationInput(
         function_inputs=FunctionInputs(
             function_input_blob=await _presign_function_input_blob(
                 data_payload=task.input,

indexify/executor/function_executor_controller/{run_task.py → run_task_allocation.py}

@@ -36,57 +36,57 @@ from indexify.proto.executor_api_pb2 import (
     TaskOutcomeCode,
 )
 
-from .events import TaskExecutionFinished
-from .metrics.run_task import (
+from .events import TaskAllocationExecutionFinished
+from .metrics.run_task_allocation import (
     metric_function_executor_run_task_rpc_errors,
     metric_function_executor_run_task_rpc_latency,
     metric_function_executor_run_task_rpcs,
     metric_function_executor_run_task_rpcs_in_progress,
 )
-from .task_info import TaskInfo
-from .task_output import TaskMetrics, TaskOutput
+from .task_allocation_info import TaskAllocationInfo
+from .task_allocation_output import TaskAllocationMetrics, TaskAllocationOutput
 
 _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
-) -> TaskExecutionFinished:
-    """Runs the task on the Function Executor and sets task_info.output with the result.
+async def run_task_allocation_on_function_executor(
+    alloc_info: TaskAllocationInfo, function_executor: FunctionExecutor, logger: Any
+) -> TaskAllocationExecutionFinished:
+    """Runs the task on the Function Executor and sets alloc_info.output with the result.
 
     Doesn't raise any exceptions.
     """
     logger = logger.bind(module=__name__)
 
-    if task_info.input is None:
+    if alloc_info.input is None:
         logger.error(
-            "task input is None, this should never happen",
+            "task allocation input is None, this should never happen",
         )
-        task_info.output = TaskOutput.internal_error(
-            allocation=task_info.allocation,
+        alloc_info.output = TaskAllocationOutput.internal_error(
+            allocation=alloc_info.allocation,
             execution_start_time=None,
             execution_end_time=None,
         )
-        return TaskExecutionFinished(
-            task_info=task_info,
+        return TaskAllocationExecutionFinished(
+            alloc_info=alloc_info,
             function_executor_termination_reason=None,
         )
 
     task = Task(
-        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,
-        request=task_info.input.function_inputs,
+        namespace=alloc_info.allocation.task.namespace,
+        graph_name=alloc_info.allocation.task.graph_name,
+        graph_version=alloc_info.allocation.task.graph_version,
+        function_name=alloc_info.allocation.task.function_name,
+        graph_invocation_id=alloc_info.allocation.task.graph_invocation_id,
+        task_id=alloc_info.allocation.task.id,
+        allocation_id=alloc_info.allocation.allocation_id,
+        request=alloc_info.input.function_inputs,
     )
 
     function_executor.invocation_state_client().add_task_to_invocation_id_entry(
-        task_id=task_info.allocation.task.id,
-        invocation_id=task_info.allocation.task.graph_invocation_id,
+        task_id=alloc_info.allocation.task.id,
+        invocation_id=alloc_info.allocation.task.graph_invocation_id,
     )
 
     metric_function_executor_run_task_rpcs.inc()
@@ -107,15 +107,15 @@ async def run_task_on_function_executor(
     # 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: float = task_info.allocation.task.timeout_ms / 1000.0
+    timeout_sec: float = alloc_info.allocation.task.timeout_ms / 1000.0
     try:
         # This aio task can only be cancelled during this await call.
         task_result = await _run_task_rpcs(task, function_executor, timeout_sec)
 
         _process_task_diagnostics(task_result.diagnostics, logger)
 
-        task_info.output = _task_output_from_function_executor_result(
-            allocation=task_info.allocation,
+        alloc_info.output = _task_alloc_output_from_fe_result(
+            allocation=alloc_info.allocation,
             result=task_result,
             execution_start_time=execution_start_time,
             execution_end_time=time.monotonic(),
@@ -127,8 +127,8 @@ async def run_task_on_function_executor(
         function_executor_termination_reason = (
             FunctionExecutorTerminationReason.FUNCTION_EXECUTOR_TERMINATION_REASON_FUNCTION_TIMEOUT
         )
-        task_info.output = TaskOutput.function_timeout(
-            allocation=task_info.allocation,
+        alloc_info.output = TaskAllocationOutput.function_timeout(
+            allocation=alloc_info.allocation,
             execution_start_time=execution_start_time,
             execution_end_time=time.monotonic(),
         )
@@ -150,14 +150,16 @@ async def run_task_on_function_executor(
             # 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)
+            logger.error(
+                "task allocationmanagement RPC execution deadline exceeded", exc_info=e
+            )
         else:
             # 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)
+            logger.error("task allocation management RPC failed", exc_info=e)
 
-        task_info.output = TaskOutput.function_executor_unresponsive(
-            allocation=task_info.allocation,
+        alloc_info.output = TaskAllocationOutput.function_executor_unresponsive(
+            allocation=alloc_info.allocation,
             execution_start_time=execution_start_time,
             execution_end_time=time.monotonic(),
         )
@@ -167,8 +169,8 @@ async def run_task_on_function_executor(
         function_executor_termination_reason = (
             FunctionExecutorTerminationReason.FUNCTION_EXECUTOR_TERMINATION_REASON_FUNCTION_CANCELLED
         )
-        task_info.output = TaskOutput.task_cancelled(
-            allocation=task_info.allocation,
+        alloc_info.output = TaskAllocationOutput.task_allocation_cancelled(
+            allocation=alloc_info.allocation,
             execution_start_time=execution_start_time,
             execution_end_time=time.monotonic(),
         )
@@ -176,10 +178,11 @@ async def run_task_on_function_executor(
         # 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
+            "unexpected internal error during task allocation lifecycle RPC sequence",
+            exc_info=e,
         )
-        task_info.output = TaskOutput.internal_error(
-            allocation=task_info.allocation,
+        alloc_info.output = TaskAllocationOutput.internal_error(
+            allocation=alloc_info.allocation,
             execution_start_time=execution_start_time,
             execution_end_time=time.monotonic(),
         )
@@ -190,11 +193,11 @@ async def run_task_on_function_executor(
     metric_function_executor_run_task_rpcs_in_progress.dec()
 
     function_executor.invocation_state_client().remove_task_to_invocation_id_entry(
-        task_id=task_info.allocation.task.id,
+        task_id=alloc_info.allocation.task.id,
     )
 
     if (
-        task_info.output.outcome_code == TaskOutcomeCode.TASK_OUTCOME_CODE_FAILURE
+        alloc_info.output.outcome_code == TaskOutcomeCode.TASK_OUTCOME_CODE_FAILURE
         and function_executor_termination_reason is None
     ):
         try:
@@ -205,7 +208,7 @@ async def run_task_on_function_executor(
                     FunctionExecutorTerminationReason.FUNCTION_EXECUTOR_TERMINATION_REASON_UNHEALTHY
                 )
                 logger.error(
-                    "Function Executor health check failed after running task, shutting down Function Executor",
+                    "Function Executor health check failed after running task allocation, shutting down Function Executor",
                     health_check_fail_reason=result.reason,
                 )
         except asyncio.CancelledError:
@@ -213,10 +216,10 @@ async def run_task_on_function_executor(
             # We can't conclude anything about the health of the FE here.
             pass
 
-    _log_task_execution_finished(output=task_info.output, logger=logger)
+    _log_task_execution_finished(output=alloc_info.output, logger=logger)
 
-    return TaskExecutionFinished(
-        task_info=task_info,
+    return TaskAllocationExecutionFinished(
+        alloc_info=alloc_info,
         function_executor_termination_reason=function_executor_termination_reason,
     )
 
@@ -275,17 +278,17 @@ async def _run_task_rpcs(
     return task_result
 
 
-def _task_output_from_function_executor_result(
+def _task_alloc_output_from_fe_result(
     allocation: TaskAllocation,
     result: TaskResult,
     execution_start_time: Optional[float],
     execution_end_time: Optional[float],
     logger: Any,
-) -> TaskOutput:
+) -> TaskAllocationOutput:
     response_validator = MessageValidator(result)
     response_validator.required_field("outcome_code")
 
-    metrics = TaskMetrics(counters={}, timers={})
+    metrics = TaskAllocationMetrics(counters={}, timers={})
     if result.HasField("metrics"):
         # Can be None if e.g. function failed.
         metrics.counters = dict(result.metrics.counters)
@@ -312,7 +315,7 @@ def _task_output_from_function_executor_result(
         # function_outputs can have no items, this happens when the function returns None.
         response_validator.required_field("uploaded_function_outputs_blob")
 
-    return TaskOutput(
+    return TaskAllocationOutput(
         allocation=allocation,
         outcome_code=outcome_code,
         failure_reason=failure_reason,
@@ -327,9 +330,9 @@ def _task_output_from_function_executor_result(
     )
 
 
-def _log_task_execution_finished(output: TaskOutput, logger: Any) -> None:
+def _log_task_execution_finished(output: TaskAllocationOutput, logger: Any) -> None:
     logger.info(
-        "finished running task",
+        "finished running task allocation",
         success=output.outcome_code == TaskOutcomeCode.TASK_OUTCOME_CODE_SUCCESS,
         outcome_code=TaskOutcomeCode.Name(output.outcome_code),
         failure_reason=(
@@ -344,7 +347,7 @@ def _process_task_diagnostics(task_diagnostics: TaskDiagnostics, logger: Any) ->
     MessageValidator(task_diagnostics).required_field("function_executor_log")
     # Uncomment these lines once we stop printing FE logs to stdout/stderr.
     # Print FE logs directly to Executor logs so operators can see them.
-    # logger.info("Function Executor logs during task execution:")
+    # logger.info("Function Executor logs during task allocation execution:")
     # print(task_diagnostics.function_executor_log)
 
 

indexify/executor/function_executor_controller/{task_info.py → task_allocation_info.py}

@@ -4,13 +4,13 @@ from typing import Optional
 
 from indexify.proto.executor_api_pb2 import TaskAllocation
 
-from .task_input import TaskInput
-from .task_output import TaskOutput
+from .task_allocation_input import TaskAllocationInput
+from .task_allocation_output import TaskAllocationOutput
 
 
 @dataclass
-class TaskInfo:
-    """Object used to track a task during its full lifecycle in the FunctionExecutorController."""
+class TaskAllocationInfo:
+    """Object used to track a task allocation during its full lifecycle in the FunctionExecutorController."""
 
     allocation: TaskAllocation
     # time.monotonic() timestamp
@@ -22,8 +22,8 @@ class TaskInfo:
     # aio task that is currently executing a lifecycle step of this task.
     aio_task: Optional[asyncio.Task] = None
     # Input if function was prepared successfully.
-    input: Optional[TaskInput] = None
+    input: Optional[TaskAllocationInput] = None
     # Output of the task, always set when the task is completed.
-    output: Optional[TaskOutput] = None
+    output: Optional[TaskAllocationOutput] = None
     # True if the task is fully completed and was added to state reporter.
     is_completed: bool = False

indexify/executor/function_executor_controller/{task_input.py → task_allocation_input.py}

@@ -1,8 +1,8 @@
 from tensorlake.function_executor.proto.function_executor_pb2 import FunctionInputs
 
 
-class TaskInput:
-    """Represents the input for a task in the function executor controller."""
+class TaskAllocationInput:
+    """Represents the input for a task allocation in the function executor controller."""
 
     def __init__(
         self,