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

indexify/executor/function_executor_controller/task_output.py

@@ -0,0 +1,122 @@
+from typing import Dict, List, Optional
+
+from tensorlake.function_executor.proto.function_executor_pb2 import (
+    FunctionOutput,
+    RouterOutput,
+)
+
+from indexify.proto.executor_api_pb2 import (
+    DataPayload,
+    Task,
+    TaskFailureReason,
+    TaskOutcomeCode,
+)
+
+
+class TaskMetrics:
+    """Metrics for a task."""
+
+    def __init__(self, counters: Dict[str, int], timers: Dict[str, float]):
+        self.counters = counters
+        self.timers = timers
+
+
+class TaskOutput:
+    """Result of running a task."""
+
+    def __init__(
+        self,
+        task: Task,
+        allocation_id: str,
+        outcome_code: TaskOutcomeCode,
+        # Optional[TaskFailureReason] is not supported in python 3.9
+        failure_reason: TaskFailureReason = None,
+        output_encoding: Optional[str] = None,
+        function_output: Optional[FunctionOutput] = None,
+        router_output: Optional[RouterOutput] = None,
+        stdout: Optional[str] = None,
+        stderr: Optional[str] = None,
+        reducer: bool = False,
+        metrics: Optional[TaskMetrics] = None,
+        uploaded_data_payloads: List[DataPayload] = [],
+        uploaded_stdout: Optional[DataPayload] = None,
+        uploaded_stderr: Optional[DataPayload] = None,
+    ):
+        self.task = task
+        self.allocation_id = allocation_id
+        self.function_output = function_output
+        self.router_output = router_output
+        self.stdout = stdout
+        self.stderr = stderr
+        self.reducer = reducer
+        self.outcome_code = outcome_code
+        self.failure_reason = failure_reason
+        self.metrics = metrics
+        self.output_encoding = output_encoding
+        self.uploaded_data_payloads = uploaded_data_payloads
+        self.uploaded_stdout = uploaded_stdout
+        self.uploaded_stderr = uploaded_stderr
+
+    @classmethod
+    def internal_error(
+        cls,
+        task: Task,
+        allocation_id: str,
+    ) -> "TaskOutput":
+        """Creates a TaskOutput for an internal error."""
+        # We are not sharing internal error messages with the customer.
+        return TaskOutput(
+            task=task,
+            allocation_id=allocation_id,
+            outcome_code=TaskOutcomeCode.TASK_OUTCOME_CODE_FAILURE,
+            failure_reason=TaskFailureReason.TASK_FAILURE_REASON_INTERNAL_ERROR,
+            stderr="Platform failed to execute the function.",
+        )
+
+    @classmethod
+    def function_timeout(
+        cls,
+        task: Task,
+        allocation_id: str,
+        timeout_sec: float,
+    ) -> "TaskOutput":
+        """Creates a TaskOutput for an function timeout error."""
+        # Task stdout, stderr is not available.
+        return TaskOutput(
+            task=task,
+            allocation_id=allocation_id,
+            outcome_code=TaskOutcomeCode.TASK_OUTCOME_CODE_FAILURE,
+            failure_reason=TaskFailureReason.TASK_FAILURE_REASON_FUNCTION_TIMEOUT,
+            stderr=f"Function exceeded its configured timeout of {timeout_sec:.3f} sec.",
+        )
+
+    @classmethod
+    def task_cancelled(
+        cls,
+        task: Task,
+        allocation_id: str,
+    ) -> "TaskOutput":
+        """Creates a TaskOutput for the case when task didn't finish because its allocation was removed by Server."""
+        return TaskOutput(
+            task=task,
+            allocation_id=allocation_id,
+            outcome_code=TaskOutcomeCode.TASK_OUTCOME_CODE_FAILURE,
+            failure_reason=TaskFailureReason.TASK_FAILURE_REASON_TASK_CANCELLED,
+        )
+
+    @classmethod
+    def function_executor_terminated(
+        cls,
+        task: Task,
+        allocation_id: str,
+    ) -> "TaskOutput":
+        """Creates a TaskOutput for the case when task didn't run because its FE terminated."""
+        return TaskOutput(
+            task=task,
+            allocation_id=allocation_id,
+            outcome_code=TaskOutcomeCode.TASK_OUTCOME_CODE_FAILURE,
+            failure_reason=TaskFailureReason.TASK_FAILURE_REASON_FUNCTION_EXECUTOR_TERMINATED,
+            # TODO: add FE startup stdout, stderr to the task output if FE failed to startup.
+            stdout="",
+            stderr="Can't execute the function because its Function Executor terminated.",
+        )

indexify/executor/function_executor_controller/upload_task_output.py

@@ -0,0 +1,234 @@
+import asyncio
+import hashlib
+import time
+from typing import Any
+
+from indexify.executor.blob_store.blob_store import BLOBStore
+from indexify.proto.executor_api_pb2 import (
+    DataPayload,
+    DataPayloadEncoding,
+)
+
+from .events import TaskOutputUploadFinished
+from .metrics.upload_task_output import (
+    metric_task_output_blob_store_upload_errors,
+    metric_task_output_blob_store_upload_latency,
+    metric_task_output_blob_store_uploads,
+    metric_task_output_upload_latency,
+    metric_task_output_upload_retries,
+    metric_task_output_uploads,
+    metric_tasks_uploading_outputs,
+)
+from .task_info import TaskInfo
+from .task_output import TaskOutput
+
+_TASK_OUTPUT_UPLOAD_BACKOFF_SEC = 5.0
+
+
+async def upload_task_output(
+    task_info: TaskInfo, blob_store: BLOBStore, logger: Any
+) -> TaskOutputUploadFinished:
+    """Uploads the task output to blob store.
+
+    Doesn't raise any Exceptions. Runs till the reporting is successful.
+    """
+    logger = logger.bind(module=__name__)
+
+    with (
+        metric_tasks_uploading_outputs.track_inprogress(),
+        metric_task_output_upload_latency.time(),
+    ):
+        metric_task_output_uploads.inc()
+        await _upload_task_output_until_successful(
+            output=task_info.output,
+            blob_store=blob_store,
+            logger=logger,
+        )
+    _log_function_metrics(output=task_info.output, logger=logger)
+    return TaskOutputUploadFinished(task_info=task_info, is_success=True)
+
+
+async def _upload_task_output_until_successful(
+    output: TaskOutput, blob_store: BLOBStore, logger: Any
+) -> None:
+    upload_retries: int = 0
+
+    while True:
+        logger = logger.bind(retries=upload_retries)
+        try:
+            await _upload_task_output_once(
+                output=output, blob_store=blob_store, logger=logger
+            )
+            return
+        except Exception as e:
+            logger.error(
+                "failed to upload task output",
+                exc_info=e,
+            )
+            upload_retries += 1
+            metric_task_output_upload_retries.inc()
+            await asyncio.sleep(_TASK_OUTPUT_UPLOAD_BACKOFF_SEC)
+
+
+class _TaskOutputSummary:
+    def __init__(self):
+        self.output_count: int = 0
+        self.output_total_bytes: int = 0
+        self.router_output_count: int = 0
+        self.stdout_count: int = 0
+        self.stdout_total_bytes: int = 0
+        self.stderr_count: int = 0
+        self.stderr_total_bytes: int = 0
+        self.total_bytes: int = 0
+
+
+async def _upload_task_output_once(
+    output: TaskOutput, blob_store: BLOBStore, logger: Any
+) -> None:
+    """Uploads the supplied task output to blob store.
+
+    Raises an Exception if the upload fails.
+    """
+    output_summary: _TaskOutputSummary = _task_output_summary(output)
+    logger.info(
+        "uploading task output to blob store",
+        total_bytes=output_summary.total_bytes,
+        total_files=output_summary.output_count
+        + output_summary.stdout_count
+        + output_summary.stderr_count,
+        output_files=output_summary.output_count,
+        output_bytes=output_summary.total_bytes,
+        router_output_count=output_summary.router_output_count,
+        stdout_bytes=output_summary.stdout_total_bytes,
+        stderr_bytes=output_summary.stderr_total_bytes,
+    )
+
+    start_time = time.time()
+    with (
+        metric_task_output_blob_store_upload_latency.time(),
+        metric_task_output_blob_store_upload_errors.count_exceptions(),
+    ):
+        metric_task_output_blob_store_uploads.inc()
+        await _upload_to_blob_store(output=output, blob_store=blob_store, logger=logger)
+
+    logger.info(
+        "files uploaded to blob store",
+        duration=time.time() - start_time,
+    )
+
+
+async def _upload_to_blob_store(
+    output: TaskOutput, blob_store: BLOBStore, logger: Any
+) -> None:
+    if output.stdout is not None:
+        stdout_url = f"{output.task.output_payload_uri_prefix}.{output.task.id}.stdout"
+        stdout_bytes: bytes = output.stdout.encode()
+        await blob_store.put(stdout_url, stdout_bytes, logger)
+        output.uploaded_stdout = DataPayload(
+            uri=stdout_url,
+            size=len(stdout_bytes),
+            sha256_hash=_compute_hash(stdout_bytes),
+            encoding=DataPayloadEncoding.DATA_PAYLOAD_ENCODING_UTF8_TEXT,
+            encoding_version=0,
+        )
+        # stdout is uploaded, free the memory used for it.
+        output.stdout = None
+
+    if output.stderr is not None:
+        stderr_url = f"{output.task.output_payload_uri_prefix}.{output.task.id}.stderr"
+        stderr_bytes: bytes = output.stderr.encode()
+        await blob_store.put(stderr_url, stderr_bytes, logger)
+        output.uploaded_stderr = DataPayload(
+            uri=stderr_url,
+            size=len(stderr_bytes),
+            sha256_hash=_compute_hash(stderr_bytes),
+            encoding=DataPayloadEncoding.DATA_PAYLOAD_ENCODING_UTF8_TEXT,
+            encoding_version=0,
+        )
+        # stderr is uploaded, free the memory used for it.
+        output.stderr = None
+
+    if output.function_output is not None:
+        # We can't use the default empty list output.uploaded_data_payloads because it's a singleton.
+        uploaded_data_payloads = []
+        for func_output_item in output.function_output.outputs:
+            node_output_sequence = len(uploaded_data_payloads)
+            output_url = f"{output.task.output_payload_uri_prefix}.{output.task.id}.{node_output_sequence}"
+            output_bytes: bytes = (
+                func_output_item.bytes
+                if func_output_item.HasField("bytes")
+                else func_output_item.string.encode()
+            )
+            await blob_store.put(output_url, output_bytes, logger)
+            uploaded_data_payloads.append(
+                DataPayload(
+                    uri=output_url,
+                    size=len(output_bytes),
+                    sha256_hash=_compute_hash(output_bytes),
+                    encoding=_to_grpc_data_payload_encoding(output),
+                    encoding_version=0,
+                )
+            )
+
+        output.uploaded_data_payloads = uploaded_data_payloads
+        # The output is uploaded, free the memory used for it.
+        output.function_output = None
+
+
+def _task_output_summary(output: TaskOutput) -> _TaskOutputSummary:
+    summary: _TaskOutputSummary = _TaskOutputSummary()
+
+    if output.stdout is not None:
+        summary.stdout_count += 1
+        summary.stdout_total_bytes += len(output.stdout)
+
+    if output.stderr is not None:
+        summary.stderr_count += 1
+        summary.stderr_total_bytes += len(output.stderr)
+
+    if output.function_output is not None:
+        for func_output_item in output.function_output.outputs:
+            output_len: bytes = len(
+                func_output_item.bytes
+                if func_output_item.HasField("bytes")
+                else func_output_item.string
+            )
+            summary.output_count += 1
+            summary.output_total_bytes += output_len
+
+    if output.router_output is not None:
+        summary.router_output_count += 1
+
+    summary.total_bytes = (
+        summary.output_total_bytes
+        + summary.stdout_total_bytes
+        + summary.stderr_total_bytes
+    )
+    return summary
+
+
+def _to_grpc_data_payload_encoding(task_output: TaskOutput) -> DataPayloadEncoding:
+    if task_output.output_encoding == "json":
+        return DataPayloadEncoding.DATA_PAYLOAD_ENCODING_UTF8_JSON
+    else:
+        return DataPayloadEncoding.DATA_PAYLOAD_ENCODING_BINARY_PICKLE
+
+
+def _compute_hash(data: bytes) -> str:
+    hasher = hashlib.sha256(usedforsecurity=False)
+    hasher.update(data)
+    return hasher.hexdigest()
+
+
+# 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
+
+    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)

indexify/executor/host_resources/host_resources.py

@@ -1,4 +1,3 @@
-import asyncio
 from typing import Any, List, Optional
 
 import psutil
@@ -47,34 +46,11 @@ class HostResourcesProvider:
             host_overhead_function_executors_ephimeral_disks_gb
         )
 
-    async def total_host_resources(self, logger: Any) -> HostResources:
+    def total_host_resources(self, logger: Any) -> HostResources:
         """Returns all hardware resources that exist at the host.
 
         Raises Exception on error.
         """
-        # Run psutil library calls in a separate thread to not block the event loop.
-        return await asyncio.to_thread(self._total_host_resources, logger=logger)
-
-    async def total_function_executor_resources(self, logger: Any) -> HostResources:
-        """Returns all hardware resources on the host that are usable by Function Executors.
-
-        Raises Exception on error.
-        """
-        total_resources: HostResources = await self.total_host_resources(logger=logger)
-        return HostResources(
-            cpu_count=max(0, total_resources.cpu_count - self._host_overhead_cpus),
-            memory_mb=max(
-                0, total_resources.memory_mb - self._host_overhead_memory_gb * 1024
-            ),
-            disk_mb=max(
-                0,
-                total_resources.disk_mb
-                - self._host_overhead_function_executors_ephimeral_disks_gb * 1024,
-            ),
-            gpus=total_resources.gpus,
-        )
-
-    def _total_host_resources(self, logger: Any) -> HostResources:
         logger = logger.bind(module=__name__)
 
         # If users disable Hyper-Threading in OS then we'd only see physical cores here.
@@ -102,3 +78,22 @@ class HostResourcesProvider:
             disk_mb=disk_mb,
             gpus=all_gpus,
         )
+
+    def total_function_executor_resources(self, logger: Any) -> HostResources:
+        """Returns all hardware resources on the host that are usable by Function Executors.
+
+        Raises Exception on error.
+        """
+        total_resources: HostResources = self.total_host_resources(logger=logger)
+        return HostResources(
+            cpu_count=max(0, total_resources.cpu_count - self._host_overhead_cpus),
+            memory_mb=max(
+                0, total_resources.memory_mb - self._host_overhead_memory_gb * 1024
+            ),
+            disk_mb=max(
+                0,
+                total_resources.disk_mb
+                - self._host_overhead_function_executors_ephimeral_disks_gb * 1024,
+            ),
+            gpus=total_resources.gpus,
+        )

indexify/executor/host_resources/nvidia_gpu_allocator.py

@@ -40,11 +40,18 @@ class NvidiaGPUAllocator:
         allocated_gpus: List[NvidiaGPUInfo] = []
         for _ in range(count):
             allocated_gpus.append(self._free_gpus.pop())
+
+        if len(allocated_gpus) > 0:
+            logger.bind(module=__name__).info("allocated GPUs:", gpus=allocated_gpus)
+
         return allocated_gpus
 
-    def deallocate(self, gpus: List[NvidiaGPUInfo]) -> None:
+    def deallocate(self, gpus: List[NvidiaGPUInfo], logger: Any) -> None:
         self._free_gpus.extend(gpus)
 
+        if len(gpus) > 0:
+            logger.bind(module=__name__).info("deallocated GPUs:", gpus=gpus)
+
     def list_all(self) -> List[NvidiaGPUInfo]:
         return list(self._all_gpus)  # Return a copy to avoid external modification
 

indexify/executor/{grpc/metrics → metrics}/channel_manager.py

@@ -1,6 +1,6 @@
 import prometheus_client
 
-from ...monitoring.metrics import latency_metric_for_fast_operation
+from ..monitoring.metrics import latency_metric_for_fast_operation
 
 metric_grpc_server_channel_creations = prometheus_client.Counter(
     "grpc_server_channel_creations",

indexify/executor/metrics/executor.py

@@ -1,10 +1,5 @@
 import prometheus_client
 
-from ..monitoring.metrics import (
-    latency_metric_for_customer_controlled_operation,
-    latency_metric_for_fast_operation,
-)
-
 # This file contains all metrics used by Executor.
 
 # Executor overview metrics.
@@ -16,45 +11,3 @@ metric_executor_state: prometheus_client.Enum = prometheus_client.Enum(
     "Current Executor state",
     states=["starting", "running", "shutting_down"],
 )
-
-# Task statistics metrics.
-metric_tasks_fetched: prometheus_client.Counter = prometheus_client.Counter(
-    "tasks_fetched", "Number of tasks that were fetched from Server"
-)
-metric_tasks_completed: prometheus_client.Counter = prometheus_client.Counter(
-    "tasks_completed", "Number of tasks that were completed", ["outcome"]
-)
-METRIC_TASKS_COMPLETED_OUTCOME_ALL = "all"
-METRIC_TASKS_COMPLETED_OUTCOME_SUCCESS = "success"
-METRIC_TASKS_COMPLETED_OUTCOME_ERROR_CUSTOMER_CODE = "error_customer_code"
-METRIC_TASKS_COMPLETED_OUTCOME_ERROR_PLATFORM = "error_platform"
-metric_tasks_completed.labels(outcome=METRIC_TASKS_COMPLETED_OUTCOME_ALL)
-metric_tasks_completed.labels(outcome=METRIC_TASKS_COMPLETED_OUTCOME_SUCCESS)
-metric_tasks_completed.labels(
-    outcome=METRIC_TASKS_COMPLETED_OUTCOME_ERROR_CUSTOMER_CODE
-)
-metric_tasks_completed.labels(outcome=METRIC_TASKS_COMPLETED_OUTCOME_ERROR_PLATFORM)
-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 outcome got reported",
-    )
-)
-
-# Task outcome reporting metrics.
-metric_task_outcome_reports: prometheus_client.Counter = prometheus_client.Counter(
-    "task_outcome_reports",
-    "Number of task outcome reports",
-)
-metric_tasks_reporting_outcome: prometheus_client.Gauge = prometheus_client.Gauge(
-    "tasks_reporting_outcome",
-    "Number of tasks currently reporting their outcomes",
-)
-metric_task_outcome_report_latency: prometheus_client.Histogram = (
-    latency_metric_for_fast_operation("task_outcome_report", "task outcome report")
-)
-metric_task_outcome_report_retries: prometheus_client.Counter = (
-    prometheus_client.Counter(
-        "tasks_outcome_report_retries", "Number of task outcome report retries"
-    )
-)

indexify/executor/{grpc/metrics → metrics}/state_reconciler.py

@@ -1,6 +1,6 @@
 import prometheus_client
 
-from ...monitoring.metrics import latency_metric_for_fast_operation
+from ..monitoring.metrics import latency_metric_for_fast_operation
 
 metric_state_reconciliations = prometheus_client.Counter(
     "state_reconciliations",

indexify/executor/{grpc/metrics → metrics}/state_reporter.py

@@ -1,6 +1,6 @@
 import prometheus_client
 
-from ...monitoring.metrics import latency_metric_for_fast_operation
+from ..monitoring.metrics import latency_metric_for_fast_operation
 
 metric_state_report_rpcs = prometheus_client.Counter(
     "state_report_rpcs",

indexify/executor/monitoring/health_checker/generic_health_checker.py

@@ -1,9 +1,3 @@
-from typing import Optional
-
-from ...function_executor.function_executor_states_container import (
-    FunctionExecutorStatesContainer,
-)
-from ...function_executor.function_executor_status import FunctionExecutorStatus
 from .health_checker import HealthChecker, HealthCheckResult
 
 HEALTH_CHECKER_NAME = "GenericHealthChecker"
@@ -16,58 +10,11 @@ class GenericHealthChecker(HealthChecker):
     """
 
     def __init__(self):
-        self._function_executor_states: Optional[FunctionExecutorStatesContainer] = None
-        self._function_executor_health_check_ever_failed = False
-
-    def set_function_executor_states_container(
-        self, states: FunctionExecutorStatesContainer
-    ):
-        self._function_executor_states = states
+        pass
 
     async def check(self) -> HealthCheckResult:
-        if self._function_executor_states is None:
-            return HealthCheckResult(
-                is_success=False,
-                status_message="Function Executor states container was not provided yet",
-                checker_name=HEALTH_CHECKER_NAME,
-            )
-
-        # Current health check policy and reasoning:
-        # * A Function Executor health check failure is a strong signal that something is wrong
-        #   either with:
-        #   - The Function Code (a criticial software bug).
-        #   - The Executor machine/container/VM (a software bug or malfunctioning local hardware).
-        # * Critical Function Code bugs tend to get fixed eventually by users. What doesn't get fixed eventually
-        #   is rare but recurring local Executor issues like hardware errors and software bugs in middleware like
-        #   drivers.
-        # * Such issues tend to get mitigated by automatically recreating the Executor machine/VM/container.
-        # * So we fail whole Executor health check if a Function Executor health check ever failed to hint the users
-        #   that we probably need to recreate the Executor machine/VM/container (unless there's a bug in Function
-        #   code that user can investigate themself).
-        await self._check_function_executors()
-        if self._function_executor_health_check_ever_failed:
-            return HealthCheckResult(
-                is_success=False,
-                status_message="A Function Executor health check failed",
-                checker_name=HEALTH_CHECKER_NAME,
-            )
-        else:
-            return HealthCheckResult(
-                is_success=True,
-                status_message="All Function Executors pass health checks",
-                checker_name=HEALTH_CHECKER_NAME,
-            )
-
-    async def _check_function_executors(self):
-        if self._function_executor_health_check_ever_failed:
-            return
-
-        async for state in self._function_executor_states:
-            # No need to async lock the state to read a single value.
-            if state.status in [
-                FunctionExecutorStatus.UNHEALTHY,
-                FunctionExecutorStatus.STARTUP_FAILED_CUSTOMER_ERROR,
-                FunctionExecutorStatus.STARTUP_FAILED_PLATFORM_ERROR,
-            ]:
-                self._function_executor_health_check_ever_failed = True
-                return
+        return HealthCheckResult(
+            is_success=True,
+            status_message="The health check is always successful",
+            checker_name=HEALTH_CHECKER_NAME,
+        )

indexify/executor/monitoring/health_checker/health_checker.py

@@ -1,8 +1,3 @@
-from ...function_executor.function_executor_states_container import (
-    FunctionExecutorStatesContainer,
-)
-
-
 class HealthCheckResult:
     def __init__(self, checker_name: str, is_success: bool, status_message: str):
         self.checker_name = checker_name
@@ -13,11 +8,5 @@ class HealthCheckResult:
 class HealthChecker:
     """Abstract base class for health checkers."""
 
-    def set_function_executor_states_container(
-        self, states: FunctionExecutorStatesContainer
-    ):
-        """Provides function executor states to this health checker so it can use them in the health checks."""
-        raise NotImplementedError("Subclasses must implement this method.")
-
     async def check(self) -> HealthCheckResult:
         raise NotImplementedError("Subclasses must implement this method.")