indexify 0.3.18__py3-none-any.whl → 0.3.19__py3-none-any.whl

indexify/executor/grpc/task_controller.py

@@ -0,0 +1,449 @@
+import asyncio
+import time
+from typing import Any, Optional
+
+import grpc
+from tensorlake.function_executor.proto.function_executor_pb2 import (
+    RunTaskRequest,
+    RunTaskResponse,
+    SerializedObject,
+)
+from tensorlake.function_executor.proto.function_executor_pb2_grpc import (
+    FunctionExecutorStub,
+)
+from tensorlake.function_executor.proto.message_validator import MessageValidator
+
+from indexify.proto.executor_api_pb2 import Task
+
+from ..downloader import Downloader
+from ..function_executor.function_executor_state import FunctionExecutorState
+from ..function_executor.function_executor_status import FunctionExecutorStatus
+from ..function_executor.metrics.single_task_runner import (
+    metric_function_executor_run_task_rpc_errors,
+    metric_function_executor_run_task_rpc_latency,
+    metric_function_executor_run_task_rpcs,
+)
+from ..function_executor.task_output import TaskMetrics, TaskOutput
+
+# TODO: combine these metrics into a single python file once gRPC migration is over and old code is removed.
+from ..metrics.executor import (
+    METRIC_TASKS_COMPLETED_OUTCOME_ALL,
+    METRIC_TASKS_COMPLETED_OUTCOME_ERROR_CUSTOMER_CODE,
+    METRIC_TASKS_COMPLETED_OUTCOME_ERROR_PLATFORM,
+    METRIC_TASKS_COMPLETED_OUTCOME_SUCCESS,
+    metric_task_completion_latency,
+    metric_task_outcome_report_latency,
+    metric_task_outcome_report_retries,
+    metric_task_outcome_reports,
+    metric_tasks_completed,
+    metric_tasks_reporting_outcome,
+)
+from ..metrics.task_runner import (
+    metric_task_policy_errors,
+    metric_task_policy_latency,
+    metric_task_policy_runs,
+    metric_task_run_latency,
+    metric_task_run_platform_errors,
+    metric_task_runs,
+    metric_tasks_blocked_by_policy,
+    metric_tasks_blocked_by_policy_per_function_name,
+    metric_tasks_running,
+)
+from ..task_reporter import TaskReporter
+from .completed_tasks_container import CompletedTasksContainer
+
+_TASK_OUTCOME_REPORT_BACKOFF_SEC = 5.0
+
+
+class FunctionTimeoutError(Exception):
+    """Exception raised when a customer's task execution exceeds the allowed timeout."""
+
+    def __init__(self, message: str):
+        super().__init__(message)
+
+
+class TaskController:
+    def __init__(
+        self,
+        task: Task,
+        function_executor_state: FunctionExecutorState,
+        downloader: Downloader,
+        task_reporter: TaskReporter,
+        completed_tasks_container: CompletedTasksContainer,
+        logger: Any,
+    ):
+        """Creates a new TaskController instance.
+
+        Raises ValueError if the supplied Task is not valid.
+        """
+        _validate_task(task)
+        self._task: Task = task
+        self._function_executor_state: FunctionExecutorState = function_executor_state
+        self._downloader: Downloader = downloader
+        self._task_reporter: TaskReporter = task_reporter
+        self._completed_tasks_container: CompletedTasksContainer = (
+            completed_tasks_container
+        )
+        self._logger: Any = logger.bind(
+            function_executor_id=function_executor_state.id,
+            task_id=task.id,
+            module=__name__,
+            namespace=task.namespace,
+            graph_name=task.graph_name,
+            graph_version=task.graph_version,
+            function_name=task.function_name,
+            invocation_id=task.graph_invocation_id,
+        )
+        self._is_running: bool = False
+        self._is_cancelled: bool = False
+        self._input: Optional[SerializedObject] = None
+        self._init_value: Optional[SerializedObject] = None
+        self._output: Optional[TaskOutput] = None
+
+    async def cancel_task(self) -> None:
+        """Cancells the task."""
+        self._is_cancelled = True
+
+        async with self._function_executor_state.lock:
+            if not self._is_running:
+                return
+
+            # Mark the Function Executor as unhealthy to destroy it to cancel the running function.
+            # If FE status changed, then it means that we're off normal task execution path, e.g.
+            # Server decided to do something with FE.
+            if (
+                self._function_executor_state.status
+                == FunctionExecutorStatus.RUNNING_TASK
+            ):
+                # TODO: Add a separate FE status for cancelled function so we don't lie to server that FE is unhealthy to destroy it.
+                await self._function_executor_state.set_status(
+                    FunctionExecutorStatus.UNHEALTHY,
+                )
+                self._logger.warning("task is cancelled")
+            else:
+                self._logger.warning(
+                    "skipping marking Function Executor unhealthy on task cancellation due to unexpected FE status",
+                    status=self._function_executor_state.status.name,
+                )
+
+    async def run_task(self) -> None:
+        """Runs the supplied task and does full managemenet of its lifecycle.
+
+        Doesn't raise any exceptions."""
+        start_time: float = time.monotonic()
+
+        try:
+            # The task can be cancelled at any time but we'll just wait until FE gets shutdown
+            # because we require this to happen from the cancel_task() caller.
+            self._input = await self._downloader.download_input(
+                namespace=self._task.namespace,
+                graph_name=self._task.graph_name,
+                graph_invocation_id=self._task.graph_invocation_id,
+                input_key=self._task.input_key,
+                logger=self._logger,
+            )
+            if self._task.HasField("reducer_output_key"):
+                self._init_value = await self._downloader.download_init_value(
+                    namespace=self._task.namespace,
+                    graph_name=self._task.graph_name,
+                    function_name=self._task.function_name,
+                    graph_invocation_id=self._task.graph_invocation_id,
+                    reducer_output_key=self._task.reducer_output_key,
+                    logger=self._logger,
+                )
+
+            await self._wait_for_idle_function_executor()
+
+            with (
+                metric_task_run_platform_errors.count_exceptions(),
+                metric_tasks_running.track_inprogress(),
+                metric_task_run_latency.time(),
+            ):
+                metric_task_runs.inc()
+                await self._run_task()
+
+            self._logger.info("task execution finished", success=self._output.success)
+        except FunctionTimeoutError:
+            self._output = TaskOutput.function_timeout(
+                task_id=self._task.id,
+                namespace=self._task.namespace,
+                graph_name=self._task.graph_name,
+                function_name=self._task.function_name,
+                graph_version=self._task.graph_version,
+                graph_invocation_id=self._task.graph_invocation_id,
+            )
+            async with self._function_executor_state.lock:
+                # Mark the Function Executor as unhealthy to destroy it to cancel the running function.
+                # If FE status changed, then it means that we're off normal task execution path, e.g.
+                # Server decided to do something with FE.
+                if (
+                    self._function_executor_state.status
+                    == FunctionExecutorStatus.RUNNING_TASK
+                ):
+                    # TODO: Add a separate FE status for timed out function so we don't lie to server that FE is unhealthy to destroy it.
+                    await self._function_executor_state.set_status(
+                        FunctionExecutorStatus.UNHEALTHY,
+                    )
+                else:
+                    self._logger.warning(
+                        "skipping marking Function Executor unhealthy on task timeout due to unexpected FE status",
+                        status=self._function_executor_state.status.name,
+                    )
+        except Exception as e:
+            self._output = TaskOutput.internal_error(
+                task_id=self._task.id,
+                namespace=self._task.namespace,
+                graph_name=self._task.graph_name,
+                function_name=self._task.function_name,
+                graph_version=self._task.graph_version,
+                graph_invocation_id=self._task.graph_invocation_id,
+            )
+            self._logger.error("task execution failed", exc_info=e)
+        finally:
+            # Release the Function Executor so others can run tasks on it if FE status didn't change.
+            # If FE status changed, then it means that we're off normal task execution path, e.g.
+            # Server decided to do something with FE.
+            async with self._function_executor_state.lock:
+                if (
+                    self._function_executor_state.status
+                    == FunctionExecutorStatus.RUNNING_TASK
+                ):
+                    await self._function_executor_state.set_status(
+                        FunctionExecutorStatus.IDLE
+                    )
+                else:
+                    self._logger.warning(
+                        "skipping marking Function Executor IDLE due to unexpected FE status",
+                        status=self._function_executor_state.status,
+                    )
+
+        _log_function_metrics(self._output, self._logger)
+
+        with (
+            metric_tasks_reporting_outcome.track_inprogress(),
+            metric_task_outcome_report_latency.time(),
+        ):
+            metric_task_outcome_reports.inc()
+            await self._report_task_outcome()
+
+        metric_task_completion_latency.observe(time.monotonic() - start_time)
+
+    async def _wait_for_idle_function_executor(self) -> None:
+        """Waits until the Function Executor is in IDLE state.
+
+        Raises an Exception if the Function Executor is in SHUTDOWN state.
+        """
+        with (
+            metric_task_policy_errors.count_exceptions(),
+            metric_tasks_blocked_by_policy.track_inprogress(),
+            metric_tasks_blocked_by_policy_per_function_name.labels(
+                function_name=self._task.function_name
+            ).track_inprogress(),
+            metric_task_policy_latency.time(),
+        ):
+            metric_task_policy_runs.inc()
+            self._logger.info(
+                "task is blocked by policy: waiting for idle function executor"
+            )
+            async with self._function_executor_state.lock:
+                await self._function_executor_state.wait_status(
+                    allowlist=[
+                        FunctionExecutorStatus.IDLE,
+                        FunctionExecutorStatus.SHUTDOWN,
+                    ]
+                )
+                if (
+                    self._function_executor_state.status
+                    == FunctionExecutorStatus.SHUTDOWN
+                ):
+                    raise Exception(
+                        "Task's Function Executor got shutdown, can't run task"
+                    )
+                await self._function_executor_state.set_status(
+                    FunctionExecutorStatus.RUNNING_TASK
+                )
+
+            # At this point the Function Executor belongs to this task controller due to RUNNING_TASK status.
+            # We can now unlock the FE state. We have to update the FE status once the task succeeds or fails.
+
+    async def _run_task(self) -> None:
+        request: RunTaskRequest = RunTaskRequest(
+            namespace=self._task.namespace,
+            graph_name=self._task.graph_name,
+            graph_version=self._task.graph_version,
+            function_name=self._task.function_name,
+            graph_invocation_id=self._task.graph_invocation_id,
+            task_id=self._task.id,
+            function_input=self._input,
+        )
+        if self._init_value is not None:
+            request.function_init_value.CopyFrom(self._init_value)
+        channel: grpc.aio.Channel = (
+            self._function_executor_state.function_executor.channel()
+        )
+
+        timeout_sec: Optional[float] = None
+        if self._task.HasField("timeout_ms"):
+            # TODO: Add integration tests with function timeout when end-to-end implementation is done.
+            timeout_sec = self._task.timeout_ms / 1000.0
+
+        async with _RunningTaskContextManager(
+            task=self._task,
+            function_executor_state=self._function_executor_state,
+        ):
+            with (
+                metric_function_executor_run_task_rpc_errors.count_exceptions(),
+                metric_function_executor_run_task_rpc_latency.time(),
+            ):
+                metric_function_executor_run_task_rpcs.inc()
+                # 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.
+                try:
+                    response: RunTaskResponse = await FunctionExecutorStub(
+                        channel
+                    ).run_task(request, timeout=timeout_sec)
+                except grpc.aio.AioRpcError as e:
+                    if e.code() == grpc.StatusCode.DEADLINE_EXCEEDED:
+                        raise FunctionTimeoutError(
+                            f"Task execution timeout {timeout_sec} expired"
+                        ) from e
+                    raise
+
+        self._output = _task_output(task=self._task, response=response)
+
+    async def _report_task_outcome(self) -> None:
+        """Reports the task with the given output to the server.
+
+        Doesn't raise any Exceptions. Runs till the reporting is successful."""
+        reporting_retries: int = 0
+
+        while True:
+            logger = self._logger.bind(retries=reporting_retries)
+            if self._is_cancelled:
+                logger.warning(
+                    "task is cancelled, skipping its outcome reporting to workaround lack of server side retries"
+                )
+                break
+
+            try:
+                await self._task_reporter.report(output=self._output, logger=logger)
+                break
+            except Exception as e:
+                logger.error(
+                    "failed to report task",
+                    exc_info=e,
+                )
+                reporting_retries += 1
+                metric_task_outcome_report_retries.inc()
+                await asyncio.sleep(_TASK_OUTCOME_REPORT_BACKOFF_SEC)
+
+        await self._completed_tasks_container.add(self._task.id)
+        metric_tasks_completed.labels(outcome=METRIC_TASKS_COMPLETED_OUTCOME_ALL).inc()
+        if self._output.is_internal_error:
+            metric_tasks_completed.labels(
+                outcome=METRIC_TASKS_COMPLETED_OUTCOME_ERROR_PLATFORM
+            ).inc()
+        elif self._output.success:
+            metric_tasks_completed.labels(
+                outcome=METRIC_TASKS_COMPLETED_OUTCOME_SUCCESS
+            ).inc()
+        else:
+            metric_tasks_completed.labels(
+                outcome=METRIC_TASKS_COMPLETED_OUTCOME_ERROR_CUSTOMER_CODE
+            ).inc()
+
+
+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("input_key")
+
+
+def _task_output(task: Task, response: RunTaskResponse) -> 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_id=task.id,
+        namespace=task.namespace,
+        graph_name=task.graph_name,
+        function_name=task.function_name,
+        graph_version=task.graph_version,
+        graph_invocation_id=task.graph_invocation_id,
+        stdout=response.stdout,
+        stderr=response.stderr,
+        reducer=response.is_reducer,
+        success=response.success,
+        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
+
+
+# Temporary workaround is logging customer metrics until we store them somewhere
+# for future retrieval and processing.
+def _log_function_metrics(output: TaskOutput, logger: Any):
+    if output.metrics is None:
+        return
+
+    logger = logger.bind(
+        invocation_id=output.graph_invocation_id,
+        function_name=output.function_name,
+        graph_name=output.graph_name,
+        namespace=output.namespace,
+    )
+
+    for counter_name, counter_value in output.metrics.counters.items():
+        logger.info(
+            "function_metric", counter_name=counter_name, counter_value=counter_value
+        )
+    for timer_name, timer_value in output.metrics.timers.items():
+        logger.info("function_metric", timer_name=timer_name, timer_value=timer_value)
+
+
+class _RunningTaskContextManager:
+    """Performs all the actions required before and after running a task."""
+
+    def __init__(
+        self,
+        task_controller: TaskController,
+    ):
+        self._task_controller: TaskController = task_controller
+
+    async def __aenter__(self):
+        self._task_controller._function_executor_state.function_executor.invocation_state_client().add_task_to_invocation_id_entry(
+            task_id=self._task_controller._task.id,
+            invocation_id=self._task_controller._task.graph_invocation_id,
+        )
+        self._task_controller._is_running = True
+        return self
+
+    async def __aexit__(self, exc_type, exc_val, exc_tb):
+        self._task_controller._is_running = False
+        self._task_controller._function_executor_state.function_executor.invocation_state_client().remove_task_to_invocation_id_entry(
+            task_id=self._task_controller._task.id,
+        )

indexify/executor/metrics/task_reporter.py

@@ -20,3 +20,17 @@ metric_server_ingest_files_latency: prometheus_client.Histogram = (
         "server_ingest_files_request", "Ingest files request to Server"
     )
 )
+
+metric_report_task_outcome_rpcs = prometheus_client.Counter(
+    "report_task_outcome_rpcs",
+    "Number of report task outcome RPCs to Server",
+)
+metric_report_task_outcome_errors = prometheus_client.Counter(
+    "report_task_outcome_rpc_errors",
+    "Number of report task outcome RPC errors",
+)
+metric_report_task_outcome_latency: prometheus_client.Histogram = (
+    latency_metric_for_fast_operation(
+        "report_task_outcome_rpc", "Report task outcome RPC to Server"
+    )
+)

indexify/executor/task_reporter.py

@@ -7,14 +7,27 @@ from httpx import Timeout
 from tensorlake.function_executor.proto.function_executor_pb2 import FunctionOutput
 from tensorlake.utils.http_client import get_httpx_client
 
+from indexify.proto.executor_api_pb2 import (
+    DataPayload,
+    OutputEncoding,
+    ReportTaskOutcomeRequest,
+    TaskOutcome,
+)
+from indexify.proto.executor_api_pb2_grpc import ExecutorAPIStub
+
 from .api_objects import (
     TASK_OUTCOME_FAILURE,
     TASK_OUTCOME_SUCCESS,
+    IngestFnOutputsResponse,
     RouterOutput,
     TaskResult,
 )
 from .function_executor.task_output import TaskOutput
+from .grpc.channel_manager import ChannelManager
 from .metrics.task_reporter import (
+    metric_report_task_outcome_errors,
+    metric_report_task_outcome_latency,
+    metric_report_task_outcome_rpcs,
     metric_server_ingest_files_errors,
     metric_server_ingest_files_latency,
     metric_server_ingest_files_requests,
@@ -45,7 +58,11 @@ class TaskOutputSummary:
 
 class TaskReporter:
     def __init__(
-        self, base_url: str, executor_id: str, config_path: Optional[str] = None
+        self,
+        base_url: str,
+        executor_id: str,
+        channel_manager: ChannelManager,
+        config_path: Optional[str] = None,
     ):
         self._base_url = base_url
         self._executor_id = executor_id
@@ -56,6 +73,7 @@ class TaskReporter:
         # Creating a new async client for each request fixes this but it
         # results in not reusing established TCP connections to server.
         self._client = get_httpx_client(config_path, make_async=False)
+        self._channel_manager = channel_manager
 
     async def shutdown(self):
         """Shuts down the task reporter.
@@ -109,12 +127,12 @@ class TaskReporter:
             # Run in a separate thread to not block the main event loop.
             response = await asyncio.to_thread(
                 self._client.post,
-                url=f"{self._base_url}/internal/ingest_files",
+                url=f"{self._base_url}/internal/ingest_fn_outputs",
                 **kwargs,
             )
         end_time = time.time()
         logger.info(
-            "task outcome reported",
+            "files uploaded",
             response_time=end_time - start_time,
             response_code=response.status_code,
         )
@@ -125,11 +143,70 @@ class TaskReporter:
             metric_server_ingest_files_errors.inc()
             # Caller catches and logs the exception.
             raise Exception(
-                "failed to report task outcome. "
+                "failed to upload files. "
                 f"Response code: {response.status_code}. "
                 f"Response text: '{response.text}'."
             ) from e
 
+        # TODO: If the files are uploaded successfully,
+        # we should record that so that if we fail to report
+        # the task outcome, we don't retry the upload.
+        # This will save us some time and resources.
+
+        ingested_files_response = response.json()
+        ingested_files = IngestFnOutputsResponse.model_validate(ingested_files_response)
+        fn_outputs = []
+        for data_payload in ingested_files.data_payloads:
+            fn_outputs.append(
+                DataPayload(
+                    path=data_payload.path,
+                    size=data_payload.size,
+                    sha256_hash=data_payload.sha256_hash,
+                )
+            )
+        stdout, stderr = None, None
+        if ingested_files.stdout:
+            stdout = DataPayload(
+                path=ingested_files.stdout.path,
+                size=ingested_files.stdout.size,
+                sha256_hash=ingested_files.stdout.sha256_hash,
+            )
+        if ingested_files.stderr:
+            stderr = DataPayload(
+                path=ingested_files.stderr.path,
+                size=ingested_files.stderr.size,
+                sha256_hash=ingested_files.stderr.sha256_hash,
+            )
+
+        request = ReportTaskOutcomeRequest(
+            task_id=output.task_id,
+            namespace=output.namespace,
+            graph_name=output.graph_name,
+            function_name=output.function_name,
+            graph_invocation_id=output.graph_invocation_id,
+            outcome=_to_grpc_task_outcome(output),
+            invocation_id=output.graph_invocation_id,
+            executor_id=self._executor_id,
+            reducer=output.reducer,
+            next_functions=(output.router_output.edges if output.router_output else []),
+            fn_outputs=fn_outputs,
+            stdout=stdout,
+            stderr=stderr,
+            output_encoding=_to_grpc_output_encoding(output),
+            output_encoding_version=0,
+        )
+        try:
+            stub = ExecutorAPIStub(await self._channel_manager.get_channel())
+            with (
+                metric_report_task_outcome_latency.time(),
+                metric_report_task_outcome_errors.count_exceptions(),
+            ):
+                metric_report_task_outcome_rpcs.inc()
+                await stub.report_task_outcome(request, timeout=5.0)
+        except Exception as e:
+            logger.error("failed to report task outcome", error=e)
+            raise e
+
     def _process_task_output(
         self, output: TaskOutput
     ) -> Tuple[TaskResult, List[Any], TaskOutputSummary]:
@@ -246,3 +323,17 @@ def _process_stderr(
     )
     summary.stderr_count += 1
     summary.stderr_total_bytes += len(stderr)
+
+
+def _to_grpc_task_outcome(task_output: TaskOutput) -> TaskOutcome:
+    if task_output.success:
+        return TaskOutcome.TASK_OUTCOME_SUCCESS
+    else:
+        return TaskOutcome.TASK_OUTCOME_FAILURE
+
+
+def _to_grpc_output_encoding(task_output: TaskOutput) -> OutputEncoding:
+    if task_output.output_encoding == "json":
+        return OutputEncoding.OUTPUT_ENCODING_JSON
+    else:
+        return OutputEncoding.OUTPUT_ENCODING_PICKLE

indexify/executor/task_runner.py

@@ -105,6 +105,7 @@ class TaskRunner:
             graph_version=task_input.task.graph_version,
             function_name=task_input.task.compute_fn,
             image_uri=task_input.task.image_uri,
+            secret_names=task_input.task.secret_names or [],
         )
         await state.lock.acquire()