indexify 0.3.31__py3-none-any.whl → 0.4.3__py3-none-any.whl

indexify/executor/function_executor_controller/loggers.py

@@ -0,0 +1,57 @@
+from typing import Any
+
+from indexify.proto.executor_api_pb2 import (
+    FunctionExecutorDescription,
+    Task,
+)
+
+
+def function_executor_logger(
+    function_executor_description: FunctionExecutorDescription, logger: Any
+) -> Any:
+    """Returns a logger bound with the FE's metadata.
+
+    The function assumes that the FE might be invalid."""
+    return logger.bind(
+        function_executor_id=(
+            function_executor_description.id
+            if function_executor_description.HasField("id")
+            else None
+        ),
+        namespace=(
+            function_executor_description.namespace
+            if function_executor_description.HasField("namespace")
+            else None
+        ),
+        graph_name=(
+            function_executor_description.graph_name
+            if function_executor_description.HasField("graph_name")
+            else None
+        ),
+        graph_version=(
+            function_executor_description.graph_version
+            if function_executor_description.HasField("graph_version")
+            else None
+        ),
+        function_name=(
+            function_executor_description.function_name
+            if function_executor_description.HasField("function_name")
+            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_name=task.graph_name if task.HasField("graph_name") else None,
+        graph_version=task.graph_version if task.HasField("graph_version") else None,
+        function_name=task.function_name if task.HasField("function_name") else None,
+        graph_invocation_id=(
+            task.graph_invocation_id if task.HasField("graph_invocation_id") else None
+        ),
+    )

indexify/executor/function_executor_controller/message_validators.py

@@ -0,0 +1,69 @@
+from tensorlake.function_executor.proto.message_validator import MessageValidator
+
+from indexify.proto.executor_api_pb2 import (
+    DataPayload,
+    FunctionExecutorDescription,
+    Task,
+)
+
+
+def validate_function_executor_description(
+    function_executor_description: FunctionExecutorDescription,
+) -> None:
+    """Validates the supplied FE description.
+
+    Raises ValueError if the description is not valid.
+    """
+    validator = MessageValidator(function_executor_description)
+    validator.required_field("id")
+    validator.required_field("namespace")
+    validator.required_field("graph_name")
+    validator.required_field("graph_version")
+    validator.required_field("function_name")
+    # image_uri is optional.
+    # secret_names can be empty.
+    validator.required_field("customer_code_timeout_ms")
+    validator.required_field("graph")
+    validator.required_field("resources")
+
+    _validate_data_payload(function_executor_description.graph)
+
+    validator = MessageValidator(function_executor_description.resources)
+    validator.required_field("cpu_ms_per_sec")
+    validator.required_field("memory_bytes")
+    validator.required_field("disk_bytes")
+
+    if function_executor_description.resources.HasField("gpu"):
+        validator = MessageValidator(function_executor_description.resources.gpu)
+        validator.required_field("count")
+        validator.required_field("model")
+
+
+def validate_task(task: Task) -> None:
+    """Validates the supplied Task.
+
+    Raises ValueError if the Task is not valid.
+    """
+    validator = MessageValidator(task)
+    validator.required_field("id")
+    validator.required_field("namespace")
+    validator.required_field("graph_name")
+    validator.required_field("graph_version")
+    validator.required_field("function_name")
+    validator.required_field("graph_invocation_id")
+    validator.required_field("timeout_ms")
+    validator.required_field("input")
+    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)
+
+
+def _validate_data_payload(data_payload: DataPayload) -> None:
+    """Validates the supplied DataPayload.
+
+    Raises ValueError if the DataPayload is not valid.
+    """
+    (MessageValidator(data_payload).required_field("uri").required_field("encoding"))

indexify/executor/function_executor_controller/metrics/completed_task_metrics.py

@@ -0,0 +1,68 @@
+import prometheus_client
+
+from indexify.executor.monitoring.metrics import (
+    latency_metric_for_customer_controlled_operation,
+)
+
+metric_tasks_completed: prometheus_client.Counter = prometheus_client.Counter(
+    "tasks_completed",
+    "Number of tasks that were completed",
+    ["outcome_code", "failure_reason"],
+)
+METRIC_TASKS_COMPLETED_OUTCOME_CODE_ALL = "all"
+METRIC_TASKS_COMPLETED_OUTCOME_CODE_SUCCESS = "success"
+METRIC_TASKS_COMPLETED_OUTCOME_CODE_FAILURE = "failure"
+
+METRIC_TASKS_COMPLETED_FAILURE_REASON_ALL = "all"
+# Used when the task is successfull.
+METRIC_TASKS_COMPLETED_FAILURE_REASON_NONE = "none"
+# Matches TASK_FAILURE_REASON_UNKNOWN
+METRIC_TASKS_COMPLETED_FAILURE_REASON_UNKNOWN = "unknown"
+# Includes all function errors including timeouts to reduce cardinality.
+METRIC_TASKS_COMPLETED_FAILURE_REASON_FUNCTION_ERROR = "function_error"
+# Includes all internal errors to reduce cardinality.
+METRIC_TASKS_COMPLETED_FAILURE_REASON_INTERNAL_ERROR = "internal_error"
+# Matches TASK_FAILURE_REASON_TASK_CANCELLED
+METRIC_TASKS_COMPLETED_FAILURE_REASON_TASK_CANCELLED = "task_cancelled"
+# Matches TASK_FAILURE_REASON_FUNCTION_EXECUTOR_TERMINATED
+METRIC_TASKS_COMPLETED_FAILURE_REASON_FUNCTION_EXECUTOR_TERMINATED = (
+    "function_executor_terminated"
+)
+
+# Valid combinations of the labels:
+metric_tasks_completed.labels(
+    outcome_code=METRIC_TASKS_COMPLETED_OUTCOME_CODE_ALL,
+    failure_reason=METRIC_TASKS_COMPLETED_FAILURE_REASON_ALL,
+)
+metric_tasks_completed.labels(
+    outcome_code=METRIC_TASKS_COMPLETED_OUTCOME_CODE_SUCCESS,
+    failure_reason=METRIC_TASKS_COMPLETED_FAILURE_REASON_NONE,
+)
+
+metric_tasks_completed.labels(
+    outcome_code=METRIC_TASKS_COMPLETED_OUTCOME_CODE_FAILURE,
+    failure_reason=METRIC_TASKS_COMPLETED_FAILURE_REASON_UNKNOWN,
+)
+metric_tasks_completed.labels(
+    outcome_code=METRIC_TASKS_COMPLETED_OUTCOME_CODE_FAILURE,
+    failure_reason=METRIC_TASKS_COMPLETED_FAILURE_REASON_FUNCTION_ERROR,
+)
+metric_tasks_completed.labels(
+    outcome_code=METRIC_TASKS_COMPLETED_OUTCOME_CODE_FAILURE,
+    failure_reason=METRIC_TASKS_COMPLETED_FAILURE_REASON_INTERNAL_ERROR,
+)
+metric_tasks_completed.labels(
+    outcome_code=METRIC_TASKS_COMPLETED_OUTCOME_CODE_FAILURE,
+    failure_reason=METRIC_TASKS_COMPLETED_FAILURE_REASON_TASK_CANCELLED,
+)
+metric_tasks_completed.labels(
+    outcome_code=METRIC_TASKS_COMPLETED_OUTCOME_CODE_FAILURE,
+    failure_reason=METRIC_TASKS_COMPLETED_FAILURE_REASON_FUNCTION_EXECUTOR_TERMINATED,
+)
+
+metric_task_completion_latency: prometheus_client.Histogram = (
+    latency_metric_for_customer_controlled_operation(
+        "task_completion",
+        "task completion from the moment it got fetched until its output got uploaded to blob store",
+    )
+)

indexify/executor/{metrics/downloader.py → function_executor_controller/metrics/downloads.py}

@@ -1,8 +1,6 @@
 import prometheus_client
 
-from ..monitoring.metrics import latency_metric_for_fast_operation
-
-# This file contains all metrics used by Downloader.
+from indexify.executor.monitoring.metrics import latency_metric_for_fast_operation
 
 # Graph download metrics
 metric_graph_downloads: prometheus_client.Counter = prometheus_client.Counter(

indexify/executor/function_executor_controller/metrics/function_executor_controller.py

@@ -0,0 +1,60 @@
+import prometheus_client
+
+from indexify.executor.monitoring.metrics import (
+    latency_metric_for_customer_controlled_operation,
+    latency_metric_for_fast_operation,
+)
+
+metric_control_loop_handle_event_latency: prometheus_client.Histogram = (
+    latency_metric_for_fast_operation(
+        "handle_function_executor_control_loop_event",
+        "Handle Function Executor control loop event",
+    )
+)
+
+metric_tasks_fetched: prometheus_client.Counter = prometheus_client.Counter(
+    "tasks_fetched", "Number of tasks that were fetched from Server"
+)
+
+metric_schedule_task_latency: prometheus_client.Histogram = (
+    latency_metric_for_customer_controlled_operation(
+        "schedule_task",
+        "Schedule a task 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_tasks_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)",
+        ["function_name"],
+    )
+)
+
+metric_function_executors_with_status: prometheus_client.Gauge = (
+    prometheus_client.Gauge(
+        "function_executors_with_status",
+        "Number of Function Executors with a particular status",
+        ["status"],
+    )
+)
+METRIC_FUNCTION_EXECUTORS_WITH_STATUS_LABEL_UNKNOWN = "unknown"
+METRIC_FUNCTION_EXECUTORS_WITH_STATUS_LABEL_PENDING = "pending"
+METRIC_FUNCTION_EXECUTORS_WITH_STATUS_LABEL_RUNNING = "running"
+METRIC_FUNCTION_EXECUTORS_WITH_STATUS_LABEL_TERMINATED = "terminated"
+
+metric_function_executors_with_status.labels(
+    status=METRIC_FUNCTION_EXECUTORS_WITH_STATUS_LABEL_UNKNOWN
+)
+metric_function_executors_with_status.labels(
+    status=METRIC_FUNCTION_EXECUTORS_WITH_STATUS_LABEL_PENDING
+)
+metric_function_executors_with_status.labels(
+    status=METRIC_FUNCTION_EXECUTORS_WITH_STATUS_LABEL_RUNNING
+)
+metric_function_executors_with_status.labels(
+    status=METRIC_FUNCTION_EXECUTORS_WITH_STATUS_LABEL_TERMINATED
+)

indexify/executor/{function_executor/metrics/single_task_runner.py → function_executor_controller/metrics/run_task.py}

@@ -1,8 +1,8 @@
 import prometheus_client
 
-from ...monitoring.metrics import latency_metric_for_customer_controlled_operation
-
-# This file contains all metrics used by SingleTaskRunner.
+from indexify.executor.monitoring.metrics import (
+    latency_metric_for_customer_controlled_operation,
+)
 
 metric_function_executor_run_task_rpcs: prometheus_client.Counter = (
     prometheus_client.Counter(
@@ -20,3 +20,9 @@ metric_function_executor_run_task_rpc_latency: prometheus_client.Histogram = (
         "function_executor_run_task_rpc", "Function Executor run task 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",
+    )
+)

indexify/executor/function_executor_controller/metrics/upload_task_output.py

@@ -0,0 +1,39 @@
+import prometheus_client
+
+from indexify.executor.monitoring.metrics import latency_metric_for_fast_operation
+
+# Task output upload metrics.
+metric_task_output_uploads: prometheus_client.Counter = prometheus_client.Counter(
+    "task_output_uploads",
+    "Number of task output uploads",
+)
+metric_tasks_uploading_outputs: prometheus_client.Gauge = prometheus_client.Gauge(
+    "tasks_uploading_output",
+    "Number of tasks currently uploading their outputs",
+)
+metric_task_output_upload_latency: prometheus_client.Histogram = (
+    latency_metric_for_fast_operation("task_output_upload", "task output upload")
+)
+metric_task_output_upload_retries: prometheus_client.Counter = (
+    prometheus_client.Counter(
+        "tasks_output_upload_retries", "Number of task output upload retries"
+    )
+)
+
+# Metrics for individual blob store operations.
+metric_task_output_blob_store_uploads: prometheus_client.Counter = (
+    prometheus_client.Counter(
+        "task_output_blob_store_uploads", "Number of task output uploads to blob store"
+    )
+)
+metric_task_output_blob_store_upload_errors: prometheus_client.Counter = (
+    prometheus_client.Counter(
+        "task_output_blob_store_upload_errors",
+        "Number of failed task output uploads to blob store",
+    )
+)
+metric_task_output_blob_store_upload_latency: prometheus_client.Histogram = (
+    latency_metric_for_fast_operation(
+        "task_output_blob_store_upload", "Upload task output to blob store"
+    )
+)

indexify/executor/function_executor_controller/prepare_task.py

@@ -0,0 +1,38 @@
+from typing import Any
+
+from indexify.executor.blob_store.blob_store import BLOBStore
+
+from .downloads import download_init_value, download_input
+from .events import TaskPreparationFinished
+from .task_info import TaskInfo
+
+
+async def prepare_task(
+    task_info: TaskInfo, blob_store: BLOBStore, logger: Any
+) -> TaskPreparationFinished:
+    """Prepares the task by downloading the input and init value if available.
+
+    Doesn't raise any exceptions.
+    """
+    logger = logger.bind(module=__name__)
+    try:
+        task_info.input = await download_input(
+            data_payload=task_info.task.input,
+            blob_store=blob_store,
+            logger=logger,
+        )
+
+        if task_info.task.HasField("reducer_input"):
+            task_info.init_value = await download_init_value(
+                data_payload=task_info.task.reducer_input,
+                blob_store=blob_store,
+                logger=logger,
+            )
+
+        return TaskPreparationFinished(task_info=task_info, is_success=True)
+    except Exception as e:
+        logger.error(
+            "Failed to prepare task",
+            exc_info=e,
+        )
+        return TaskPreparationFinished(task_info=task_info, is_success=False)

indexify/executor/function_executor_controller/run_task.py

@@ -0,0 +1,201 @@
+import asyncio
+import time
+from typing import Any, Optional
+
+import grpc
+from tensorlake.function_executor.proto.function_executor_pb2 import (
+    RunTaskRequest,
+    RunTaskResponse,
+)
+from tensorlake.function_executor.proto.function_executor_pb2_grpc import (
+    FunctionExecutorStub,
+)
+from tensorlake.function_executor.proto.message_validator import MessageValidator
+
+from indexify.executor.function_executor.function_executor import FunctionExecutor
+from indexify.executor.function_executor.health_checker import HealthCheckResult
+from indexify.proto.executor_api_pb2 import (
+    FunctionExecutorTerminationReason,
+    Task,
+    TaskFailureReason,
+    TaskOutcomeCode,
+)
+
+from .events import TaskExecutionFinished
+from .metrics.run_task 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
+
+
+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.
+
+    Doesn't raise any exceptions.
+    """
+    logger = logger.bind(module=__name__)
+    request: RunTaskRequest = RunTaskRequest(
+        namespace=task_info.task.namespace,
+        graph_name=task_info.task.graph_name,
+        graph_version=task_info.task.graph_version,
+        function_name=task_info.task.function_name,
+        graph_invocation_id=task_info.task.graph_invocation_id,
+        task_id=task_info.task.id,
+        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)
+        # Don't keep the init value in memory after we started running the task.
+        task_info.init_value = None
+
+    function_executor.invocation_state_client().add_task_to_invocation_id_entry(
+        task_id=task_info.task.id,
+        invocation_id=task_info.task.graph_invocation_id,
+    )
+
+    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
+
+    # 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.task.timeout_ms / 1000.0
+    try:
+        channel: grpc.aio.Channel = function_executor.channel()
+        response: RunTaskResponse = await FunctionExecutorStub(channel).run_task(
+            request, timeout=timeout_sec
+        )
+        task_info.output = _task_output_from_function_executor_response(
+            task=task_info.task,
+            response=response,
+            allocation_id=task_info.allocation_id,
+        )
+    except grpc.aio.AioRpcError as e:
+        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(
+                task=task_info.task,
+                allocation_id=task_info.allocation_id,
+                timeout_sec=timeout_sec,
+            )
+        else:
+            metric_function_executor_run_task_rpc_errors.inc()
+            logger.error("task execution failed", exc_info=e)
+            task_info.output = TaskOutput.internal_error(
+                task=task_info.task, allocation_id=task_info.allocation_id
+            )
+    except asyncio.CancelledError:
+        # The task is still running in FE, we only cancelled the client-side RPC.
+        function_executor_termination_reason = (
+            FunctionExecutorTerminationReason.FUNCTION_EXECUTOR_TERMINATION_REASON_FUNCTION_CANCELLED
+        )
+        task_info.output = TaskOutput.task_cancelled(
+            task=task_info.task, allocation_id=task_info.allocation_id
+        )
+    except Exception as e:
+        metric_function_executor_run_task_rpc_errors.inc()
+        logger.error("task execution failed", exc_info=e)
+        task_info.output = TaskOutput.internal_error(
+            task=task_info.task, allocation_id=task_info.allocation_id
+        )
+
+    metric_function_executor_run_task_rpc_latency.observe(time.monotonic() - start_time)
+    metric_function_executor_run_task_rpcs_in_progress.dec()
+
+    function_executor.invocation_state_client().remove_task_to_invocation_id_entry(
+        task_id=task_info.task.id,
+    )
+
+    if (
+        task_info.output.outcome_code == TaskOutcomeCode.TASK_OUTCOME_CODE_FAILURE
+        and function_executor_termination_reason is None
+    ):
+        # Check if the task failed because the FE is unhealthy to prevent more tasks failing.
+        result: HealthCheckResult = await function_executor.health_checker().check()
+        if not result.is_healthy:
+            function_executor_termination_reason = (
+                FunctionExecutorTerminationReason.FUNCTION_EXECUTOR_TERMINATION_REASON_UNHEALTHY
+            )
+            logger.error(
+                "Function Executor health check failed after running task, shutting down Function Executor",
+                health_check_fail_reason=result.reason,
+            )
+
+    _log_task_execution_finished(output=task_info.output, logger=logger)
+
+    return TaskExecutionFinished(
+        task_info=task_info,
+        function_executor_termination_reason=function_executor_termination_reason,
+    )
+
+
+def _task_output_from_function_executor_response(
+    task: Task, response: RunTaskResponse, allocation_id: str
+) -> TaskOutput:
+    response_validator = MessageValidator(response)
+    response_validator.required_field("stdout")
+    response_validator.required_field("stderr")
+    response_validator.required_field("is_reducer")
+    response_validator.required_field("success")
+
+    metrics = TaskMetrics(counters={}, timers={})
+    if response.HasField("metrics"):
+        # Can be None if e.g. function failed.
+        metrics.counters = dict(response.metrics.counters)
+        metrics.timers = dict(response.metrics.timers)
+
+    output = TaskOutput(
+        task=task,
+        allocation_id=allocation_id,
+        outcome_code=(
+            TaskOutcomeCode.TASK_OUTCOME_CODE_SUCCESS
+            if response.success
+            else TaskOutcomeCode.TASK_OUTCOME_CODE_FAILURE
+        ),
+        failure_reason=(
+            None
+            if response.success
+            else TaskFailureReason.TASK_FAILURE_REASON_FUNCTION_ERROR
+        ),
+        stdout=response.stdout,
+        stderr=response.stderr,
+        reducer=response.is_reducer,
+        metrics=metrics,
+    )
+
+    if response.HasField("function_output"):
+        output.function_output = response.function_output
+    if response.HasField("router_output"):
+        output.router_output = response.router_output
+
+    return output
+
+
+def _log_task_execution_finished(output: TaskOutput, logger: Any) -> None:
+    logger.info(
+        "finished running task",
+        success=output.outcome_code == TaskOutcomeCode.TASK_OUTCOME_CODE_SUCCESS,
+        outcome_code=TaskOutcomeCode.Name(output.outcome_code),
+        failure_reason=(
+            TaskFailureReason.Name(output.failure_reason)
+            if output.failure_reason is not None
+            else None
+        ),
+    )

indexify/executor/function_executor_controller/task_info.py

@@ -0,0 +1,33 @@
+import asyncio
+from dataclasses import dataclass
+from typing import Optional
+
+from tensorlake.function_executor.proto.function_executor_pb2 import SerializedObject
+
+from indexify.proto.executor_api_pb2 import Task
+
+from .task_output import TaskOutput
+
+
+@dataclass
+class TaskInfo:
+    """Object used to track a task during its full lifecycle in the FunctionExecutorController."""
+
+    task: Task
+    allocation_id: str
+    # time.monotonic() timestamp
+    start_time: float
+    # time.monotonic() timestamp when the task was prepared for execution
+    prepared_time: float = 0.0
+    # True if the task was cancelled.
+    is_cancelled: bool = False
+    # aio task that is currently executing a lifecycle step of this task.
+    aio_task: Optional[asyncio.Task] = None
+    # Downloaded input if function was prepared successfully.
+    input: Optional[SerializedObject] = None
+    # Downloaded init value if function was prepared successfully and is a reducer.
+    init_value: Optional[SerializedObject] = None
+    # Output of the task.
+    output: Optional[TaskOutput] = None
+    # True if the task is fully completed and was added to state reporter.
+    is_completed: bool = False