indexify 0.3.30__py3-none-any.whl → 0.4.2__py3-none-any.whl

indexify/executor/function_executor/function_executor_status.py

@@ -1,95 +0,0 @@
-from enum import Enum
-
-
-class FunctionExecutorStatus(Enum):
-    """Status of a Function Executor.
-
-    Each status lists transitions allowed to it.
-    """
-
-    # DESTROYED -> STARTING_UP
-    STARTING_UP = "Starting Up"
-    # STARTING_UP -> STARTUP_FAILED_CUSTOMER_ERROR
-    STARTUP_FAILED_CUSTOMER_ERROR = "Startup Failed (Customer Error)"
-    # STARTING_UP -> STARTUP_FAILED_PLATFORM_ERROR
-    STARTUP_FAILED_PLATFORM_ERROR = "Startup Failed (Platform Error)"
-    # STARTING_UP -> IDLE
-    # RUNNING_TASK -> IDLE
-    IDLE = "Idle"
-    # IDLE -> RUNNING_TASK
-    RUNNING_TASK = "Running Task"
-    # IDLE -> UNHEALTHY
-    # RUNNING_TASK -> UNHEALTHY
-    UNHEALTHY = "Unhealthy"
-    # STARTUP_FAILED_CUSTOMER_ERROR -> DESTROYING
-    # STARTUP_FAILED_PLATFORM_ERROR -> DESTROYING
-    # RUNNING_TASK -> DESTROYING
-    # UNHEALTHY -> DESTROYING
-    # IDLE -> DESTROYING
-    DESTROYING = "Destroying"
-    # DESTROYED (initial status)
-    # DESTROYING -> DESTROYED
-    DESTROYED = "Destroyed"
-    # Any state -> SHUTDOWN
-    SHUTDOWN = "Shutdown"  # Permanent stop state
-
-
-# TODO: After removing HTTP code simplify state transitions by not allowing to
-# startup an FE after it was destroyed. grpc protocol treats FEs as ephimeral and never revives them.
-def is_status_change_allowed(
-    current_status: FunctionExecutorStatus, new_status: FunctionExecutorStatus
-) -> bool:
-    """Returns True if the transition is allowed."""
-    allowed_transitions = {
-        FunctionExecutorStatus.DESTROYED: [
-            FunctionExecutorStatus.DESTROYED,
-            FunctionExecutorStatus.STARTING_UP,
-            FunctionExecutorStatus.SHUTDOWN,
-        ],
-        FunctionExecutorStatus.STARTING_UP: [
-            FunctionExecutorStatus.STARTING_UP,
-            FunctionExecutorStatus.STARTUP_FAILED_CUSTOMER_ERROR,
-            FunctionExecutorStatus.STARTUP_FAILED_PLATFORM_ERROR,
-            FunctionExecutorStatus.IDLE,
-            FunctionExecutorStatus.SHUTDOWN,
-        ],
-        FunctionExecutorStatus.STARTUP_FAILED_CUSTOMER_ERROR: [
-            FunctionExecutorStatus.STARTUP_FAILED_CUSTOMER_ERROR,
-            FunctionExecutorStatus.DESTROYING,
-            FunctionExecutorStatus.SHUTDOWN,
-        ],
-        FunctionExecutorStatus.STARTUP_FAILED_PLATFORM_ERROR: [
-            FunctionExecutorStatus.STARTUP_FAILED_PLATFORM_ERROR,
-            FunctionExecutorStatus.DESTROYING,
-            FunctionExecutorStatus.SHUTDOWN,
-        ],
-        FunctionExecutorStatus.IDLE: [
-            FunctionExecutorStatus.IDLE,
-            FunctionExecutorStatus.RUNNING_TASK,
-            FunctionExecutorStatus.UNHEALTHY,
-            FunctionExecutorStatus.DESTROYING,
-            FunctionExecutorStatus.SHUTDOWN,
-        ],
-        FunctionExecutorStatus.RUNNING_TASK: [
-            FunctionExecutorStatus.RUNNING_TASK,
-            FunctionExecutorStatus.DESTROYING,
-            FunctionExecutorStatus.IDLE,
-            FunctionExecutorStatus.UNHEALTHY,
-            FunctionExecutorStatus.SHUTDOWN,
-        ],
-        FunctionExecutorStatus.UNHEALTHY: [
-            FunctionExecutorStatus.UNHEALTHY,
-            FunctionExecutorStatus.DESTROYING,
-            FunctionExecutorStatus.SHUTDOWN,
-        ],
-        FunctionExecutorStatus.DESTROYING: [
-            FunctionExecutorStatus.DESTROYING,
-            FunctionExecutorStatus.DESTROYED,
-            FunctionExecutorStatus.SHUTDOWN,
-        ],
-        FunctionExecutorStatus.SHUTDOWN: [
-            FunctionExecutorStatus.SHUTDOWN
-        ],  # No transitions allowed from SHUTDOWN
-    }
-
-    return new_status in allowed_transitions.get(current_status, [])

indexify/executor/function_executor/metrics/function_executor_state.py

@@ -1,46 +0,0 @@
-import prometheus_client
-
-from ..function_executor_status import FunctionExecutorStatus
-
-# This file contains all metrics used by FunctionExecutorState.
-
-metric_function_executor_state_not_locked_errors: prometheus_client.Counter = (
-    prometheus_client.Counter(
-        "function_executor_state_not_locked_errors",
-        "Number of times a Function Executor state was used without acquiring its lock",
-    )
-)
-
-# Function Executors count with a particular status.
-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.labels(
-    status=FunctionExecutorStatus.STARTING_UP.name
-)
-metric_function_executors_with_status.labels(
-    status=FunctionExecutorStatus.STARTUP_FAILED_CUSTOMER_ERROR.name
-)
-metric_function_executors_with_status.labels(
-    status=FunctionExecutorStatus.STARTUP_FAILED_PLATFORM_ERROR.name
-)
-metric_function_executors_with_status.labels(status=FunctionExecutorStatus.IDLE.name)
-metric_function_executors_with_status.labels(
-    status=FunctionExecutorStatus.RUNNING_TASK.name
-)
-metric_function_executors_with_status.labels(
-    status=FunctionExecutorStatus.UNHEALTHY.name
-)
-metric_function_executors_with_status.labels(
-    status=FunctionExecutorStatus.DESTROYING.name
-)
-metric_function_executors_with_status.labels(
-    status=FunctionExecutorStatus.DESTROYED.name
-)
-metric_function_executors_with_status.labels(
-    status=FunctionExecutorStatus.SHUTDOWN.name
-)

indexify/executor/function_executor/metrics/function_executor_state_container.py

@@ -1,10 +0,0 @@
-import prometheus_client
-
-# This file contains all metrics used by FunctionExecutorStatesContainer.
-
-metric_function_executor_states_count: prometheus_client.Gauge = (
-    prometheus_client.Gauge(
-        "function_executor_states_count",
-        "Number of existing Function Executor states",
-    )
-)

indexify/executor/function_executor/single_task_runner.py

@@ -1,345 +0,0 @@
-from collections.abc import Awaitable, Callable
-from math import ceil
-from typing import Any, Optional
-
-import grpc
-from tensorlake.function_executor.proto.function_executor_pb2 import (
-    InitializeRequest,
-    RunTaskRequest,
-    RunTaskResponse,
-)
-from tensorlake.function_executor.proto.function_executor_pb2_grpc import (
-    FunctionExecutorStub,
-)
-from tensorlake.function_executor.proto.message_validator import MessageValidator
-
-from ..api_objects import Task
-from .function_executor import CustomerError, FunctionExecutor
-from .function_executor_state import FunctionExecutorState
-from .function_executor_status import FunctionExecutorStatus
-from .health_checker import HealthChecker, HealthCheckResult
-from .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 .server.function_executor_server_factory import (
-    FunctionExecutorServerConfiguration,
-    FunctionExecutorServerFactory,
-)
-from .task_input import TaskInput
-from .task_output import TaskMetrics, TaskOutput
-
-
-class SingleTaskRunner:
-    def __init__(
-        self,
-        executor_id: str,
-        function_executor_state: FunctionExecutorState,
-        task_input: TaskInput,
-        function_executor_server_factory: FunctionExecutorServerFactory,
-        base_url: str,
-        config_path: Optional[str],
-        logger: Any,
-    ):
-        self._executor_id: str = executor_id
-        self._function_executor_state: FunctionExecutorState = function_executor_state
-        self._task_input: TaskInput = task_input
-        self._function_executor_server_factory: FunctionExecutorServerFactory = (
-            function_executor_server_factory
-        )
-        self._base_url: str = base_url
-        self._config_path: Optional[str] = config_path
-        self._logger = logger.bind(module=__name__)
-
-    async def run(self) -> TaskOutput:
-        """Runs the task in the Function Executor.
-
-        The FunctionExecutorState must be locked by the caller.
-        The lock is released during actual task run in the server.
-        The lock is relocked on return.
-
-        Raises an exception if an error occured.
-
-        On enter the Function Executor status is either IDLE, UNHEALTHY or DESTROYED.
-        On return the Function Executor status is either IDLE, UNHEALTHY or DESTROYED.
-        """
-        self._function_executor_state.check_locked()
-
-        if self._function_executor_state.status not in [
-            FunctionExecutorStatus.IDLE,
-            FunctionExecutorStatus.UNHEALTHY,
-            FunctionExecutorStatus.DESTROYED,
-        ]:
-            self._logger.error(
-                "Function Executor is not in oneof [IDLE, UNHEALTHY, DESTROYED] state, cannot run the task",
-                status=self._function_executor_state.status,
-            )
-            raise RuntimeError(
-                f"Unexpected Function Executor state {self._function_executor_state.status}"
-            )
-
-        # If Function Executor became unhealthy while was idle then destroy it.
-        # It'll be recreated below.
-        await self._destroy_existing_function_executor_if_unhealthy()
-
-        # Create Function Executor if it doesn't exist yet.
-        if self._function_executor_state.status == FunctionExecutorStatus.DESTROYED:
-            try:
-                await self._create_function_executor()
-            except CustomerError as e:
-                return TaskOutput(
-                    task_id=self._task_input.task.id,
-                    namespace=self._task_input.task.namespace,
-                    graph_name=self._task_input.task.compute_graph,
-                    function_name=self._task_input.task.compute_fn,
-                    graph_version=self._task_input.task.graph_version,
-                    graph_invocation_id=self._task_input.task.invocation_id,
-                    stderr=str(e),
-                    success=False,
-                    output_payload_uri_prefix=self._task_input.task.output_payload_uri_prefix,
-                )
-
-        try:
-            return await self._run()
-        finally:
-            # If Function Executor became unhealthy while running the task then destroy it.
-            # The periodic health checker might not notice this as it does only periodic checks.
-            await self._destroy_existing_function_executor_if_unhealthy()
-
-            if self._function_executor_state.status not in [
-                FunctionExecutorStatus.IDLE,
-                FunctionExecutorStatus.UNHEALTHY,
-                FunctionExecutorStatus.DESTROYED,
-            ]:
-                self._logger.error(
-                    "Function Executor status is not oneof [IDLE, UNHEALTHY, DESTROYED] after running the task, resetting the state to mitigate a possible bug",
-                    status=self._function_executor_state.status,
-                )
-                if self._function_executor_state.function_executor is None:
-                    await self._function_executor_state.set_status(
-                        FunctionExecutorStatus.DESTROYED
-                    )
-                else:
-                    await self._function_executor_state.set_status(
-                        FunctionExecutorStatus.UNHEALTHY
-                    )
-
-    async def _create_function_executor(self) -> None:
-        await self._function_executor_state.set_status(
-            FunctionExecutorStatus.STARTING_UP
-        )
-        self._function_executor_state.function_executor = FunctionExecutor(
-            server_factory=self._function_executor_server_factory, logger=self._logger
-        )
-        task: Task = self._task_input.task
-        config: FunctionExecutorServerConfiguration = (
-            FunctionExecutorServerConfiguration(
-                executor_id=self._executor_id,
-                function_executor_id=self._function_executor_state.id,
-                namespace=task.namespace,
-                graph_name=task.compute_graph,
-                graph_version=task.graph_version,
-                function_name=task.compute_fn,
-                image_uri=task.image_uri,
-                secret_names=task.secret_names or [],
-                cpu_ms_per_sec=(
-                    None
-                    if task.resources.cpus is None
-                    else ceil(task.resources.cpus * 1000)
-                ),
-                memory_bytes=(
-                    None
-                    if task.resources.memory_mb is None
-                    else task.resources.memory_mb * 1024 * 1024
-                ),
-                disk_bytes=(
-                    None
-                    if task.resources.ephemeral_disk_mb is None
-                    else task.resources.ephemeral_disk_mb * 1024 * 1024
-                ),
-                gpu_count=0 if task.resources.gpu is None else task.resources.gpu.count,
-            )
-        )
-        initialize_request: InitializeRequest = InitializeRequest(
-            namespace=self._task_input.task.namespace,
-            graph_name=self._task_input.task.compute_graph,
-            graph_version=self._task_input.task.graph_version,
-            function_name=self._task_input.task.compute_fn,
-            graph=self._task_input.graph,
-        )
-
-        try:
-            await self._function_executor_state.function_executor.initialize(
-                config=config,
-                initialize_request=initialize_request,
-                base_url=self._base_url,
-                config_path=self._config_path,
-            )
-        except CustomerError:
-            # We have to follow the valid state transition sequence.
-            await self._function_executor_state.set_status(
-                FunctionExecutorStatus.STARTUP_FAILED_CUSTOMER_ERROR
-            )
-            await self._function_executor_state.destroy_function_executor()
-            raise
-        except Exception:
-            # We have to follow the valid state transition sequence.
-            await self._function_executor_state.set_status(
-                FunctionExecutorStatus.STARTUP_FAILED_PLATFORM_ERROR
-            )
-            await self._function_executor_state.destroy_function_executor()
-            raise
-
-        await self._function_executor_state.set_status(FunctionExecutorStatus.IDLE)
-
-    async def _run(self) -> TaskOutput:
-        request: RunTaskRequest = RunTaskRequest(
-            namespace=self._task_input.task.namespace,
-            graph_name=self._task_input.task.compute_graph,
-            graph_version=self._task_input.task.graph_version,
-            function_name=self._task_input.task.compute_fn,
-            graph_invocation_id=self._task_input.task.invocation_id,
-            task_id=self._task_input.task.id,
-            function_input=self._task_input.input,
-        )
-        if self._task_input.init_value is not None:
-            request.function_init_value.CopyFrom(self._task_input.init_value)
-        channel: grpc.aio.Channel = (
-            self._function_executor_state.function_executor.channel()
-        )
-
-        async with _RunningTaskContextManager(
-            invocation_id=self._task_input.task.invocation_id,
-            task_id=self._task_input.task.id,
-            health_check_failed_callback=self._health_check_failed_callback,
-            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.
-                response: RunTaskResponse = await FunctionExecutorStub(
-                    channel
-                ).run_task(request)
-            return _task_output(task=self._task_input.task, response=response)
-
-    async def _health_check_failed_callback(self, result: HealthCheckResult):
-        # Function Executor destroy due to the periodic health check failure ensures that
-        # a running task RPC stuck in unhealthy Function Executor fails immidiately.
-        async with self._function_executor_state.lock:
-            if (
-                self._function_executor_state.status
-                != FunctionExecutorStatus.RUNNING_TASK
-            ):
-                # Protection in case the callback gets delivered after we finished running the task.
-                return
-
-            await self._function_executor_state.set_status(
-                FunctionExecutorStatus.UNHEALTHY
-            )
-            await self._destroy_function_executor_on_failed_health_check(result.reason)
-
-    async def _destroy_existing_function_executor_if_unhealthy(self):
-        self._function_executor_state.check_locked()
-        if self._function_executor_state.status == FunctionExecutorStatus.IDLE:
-            result: HealthCheckResult = (
-                await self._function_executor_state.function_executor.health_checker().check()
-            )
-            if not result.is_healthy:
-                await self._function_executor_state.set_status(
-                    FunctionExecutorStatus.UNHEALTHY
-                )
-
-        if self._function_executor_state.status == FunctionExecutorStatus.UNHEALTHY:
-            await self._destroy_function_executor_on_failed_health_check(result.reason)
-
-    async def _destroy_function_executor_on_failed_health_check(self, reason: str):
-        self._function_executor_state.check_locked()
-        self._logger.error(
-            "Function Executor health check failed, destroying Function Executor",
-            health_check_fail_reason=reason,
-        )
-        await self._function_executor_state.destroy_function_executor()
-
-
-class _RunningTaskContextManager:
-    """Performs all the actions required before and after running a task."""
-
-    def __init__(
-        self,
-        invocation_id: str,
-        task_id: str,
-        health_check_failed_callback: Callable[[], Awaitable[None]],
-        function_executor_state: FunctionExecutorState,
-    ):
-        self._invocation_id: str = invocation_id
-        self._task_id: str = task_id
-        self._health_check_failed_callback: Callable[[], Awaitable[None]] = (
-            health_check_failed_callback
-        )
-        self._state: FunctionExecutorState = function_executor_state
-
-    async def __aenter__(self):
-        await self._state.set_status(FunctionExecutorStatus.RUNNING_TASK)
-        self._state.function_executor.invocation_state_client().add_task_to_invocation_id_entry(
-            task_id=self._task_id,
-            invocation_id=self._invocation_id,
-        )
-        self._state.function_executor.health_checker().start(
-            self._health_check_failed_callback
-        )
-        # Unlock the state so other tasks can act depending on it.
-        self._state.lock.release()
-        return self
-
-    async def __aexit__(self, exc_type, exc_val, exc_tb):
-        await self._state.lock.acquire()
-        # Health check callback could destroy the FunctionExecutor and set status to UNHEALTHY
-        if self._state.status == FunctionExecutorStatus.RUNNING_TASK:
-            await self._state.set_status(FunctionExecutorStatus.IDLE)
-            self._state.function_executor.invocation_state_client().remove_task_to_invocation_id_entry(
-                task_id=self._task_id
-            )
-            self._state.function_executor.health_checker().stop()
-
-
-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.compute_graph,
-        function_name=task.compute_fn,
-        graph_version=task.graph_version,
-        graph_invocation_id=task.invocation_id,
-        stdout=response.stdout,
-        stderr=response.stderr,
-        reducer=response.is_reducer,
-        success=response.success,
-        metrics=metrics,
-        output_payload_uri_prefix=task.output_payload_uri_prefix,
-    )
-
-    if response.HasField("function_output"):
-        output.function_output = response.function_output
-        output.output_encoding = response.function_output.output_encoding
-    if response.HasField("router_output"):
-        output.router_output = response.router_output
-
-    return output

indexify/executor/function_executor/task_input.py

@@ -1,21 +0,0 @@
-from typing import Optional
-
-from tensorlake.function_executor.proto.function_executor_pb2 import SerializedObject
-
-from ..api_objects import Task
-
-
-class TaskInput:
-    """Task with all the resources required to run it."""
-
-    def __init__(
-        self,
-        task: Task,
-        graph: SerializedObject,
-        input: SerializedObject,
-        init_value: Optional[SerializedObject],
-    ):
-        self.task: Task = task
-        self.graph: SerializedObject = graph
-        self.input: SerializedObject = input
-        self.init_value: Optional[SerializedObject] = init_value

indexify/executor/function_executor/task_output.py

@@ -1,105 +0,0 @@
-from typing import Dict, Optional
-
-from tensorlake.function_executor.proto.function_executor_pb2 import (
-    FunctionOutput,
-    RouterOutput,
-)
-
-
-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_id: str,
-        namespace: str,
-        graph_name: str,
-        function_name: str,
-        graph_version: str,
-        graph_invocation_id: str,
-        output_payload_uri_prefix: Optional[str],
-        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,
-        success: bool = False,
-        is_internal_error: bool = False,
-        metrics: Optional[TaskMetrics] = None,
-    ):
-        self.task_id = task_id
-        self.namespace = namespace
-        self.graph_name = graph_name
-        self.function_name = function_name
-        self.graph_version = graph_version
-        self.graph_invocation_id = graph_invocation_id
-        self.function_output = function_output
-        self.router_output = router_output
-        self.stdout = stdout
-        self.stderr = stderr
-        self.reducer = reducer
-        self.success = success
-        self.is_internal_error = is_internal_error
-        self.metrics = metrics
-        self.output_encoding = output_encoding
-        self.output_payload_uri_prefix = output_payload_uri_prefix
-
-    @classmethod
-    def internal_error(
-        cls,
-        task_id: str,
-        namespace: str,
-        graph_name: str,
-        function_name: str,
-        graph_version: str,
-        graph_invocation_id: str,
-        output_payload_uri_prefix: Optional[str],
-    ) -> "TaskOutput":
-        """Creates a TaskOutput for an internal error."""
-        # We are not sharing internal error messages with the customer.
-        return TaskOutput(
-            task_id=task_id,
-            namespace=namespace,
-            graph_name=graph_name,
-            function_name=function_name,
-            graph_version=graph_version,
-            graph_invocation_id=graph_invocation_id,
-            stderr="Platform failed to execute the function.",
-            is_internal_error=True,
-            output_payload_uri_prefix=output_payload_uri_prefix,
-        )
-
-    @classmethod
-    def function_timeout(
-        cls,
-        task_id: str,
-        namespace: str,
-        graph_name: str,
-        function_name: str,
-        graph_version: str,
-        graph_invocation_id: str,
-        timeout_sec: float,
-        output_payload_uri_prefix: Optional[str],
-    ) -> "TaskOutput":
-        """Creates a TaskOutput for an function timeout error."""
-        # Task stdout, stderr is not available.
-        return TaskOutput(
-            task_id=task_id,
-            namespace=namespace,
-            graph_name=graph_name,
-            function_name=function_name,
-            graph_version=graph_version,
-            graph_invocation_id=graph_invocation_id,
-            stderr=f"Function or router exceeded its configured timeout of {timeout_sec:.3f} sec.",
-            is_internal_error=False,
-            output_payload_uri_prefix=output_payload_uri_prefix,
-        )