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

indexify/executor/function_executor/server/subprocess_function_executor_server_factory.py

@@ -11,15 +11,11 @@ from .function_executor_server_factory import (
 from .subprocess_function_executor_server import SubprocessFunctionExecutorServer
 
 
-def get_free_tcp_port(iface_name="localhost") -> int:
-    tcp = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
-    tcp.bind((iface_name, 0))
-    _, port = tcp.getsockname()
-    tcp.close()
-    return port
-
-
 class SubprocessFunctionExecutorServerFactory(FunctionExecutorServerFactory):
+    def __init__(self, verbose_logs: bool) -> None:
+        super().__init__()
+        self._verbose_logs = verbose_logs
+
     async def create(
         self, config: FunctionExecutorServerConfiguration, logger: Any
     ) -> SubprocessFunctionExecutorServer:
@@ -33,13 +29,15 @@ class SubprocessFunctionExecutorServerFactory(FunctionExecutorServerFactory):
             )
 
         try:
-            port = get_free_tcp_port()
-            logger.info("allocated function executor port", port=port)
+            port = _find_free_localhost_tcp_port()
             args = [
                 f"--executor-id={config.executor_id}",  # use = as executor_id can start with -
+                f"--function-executor-id={config.function_executor_id}",
                 "--address",
                 _server_address(port),
             ]
+            if self._verbose_logs:
+                args.append("--dev")
             # Run the process with our stdout, stderr. We want to see process logs and exceptions in our process output.
             # This is useful for dubugging. Customer function stdout and stderr is captured and returned in the response
             # so we won't see it in our process outputs. This is the right behavior as customer function stdout and stderr
@@ -91,10 +89,23 @@ class SubprocessFunctionExecutorServerFactory(FunctionExecutorServerFactory):
             )
 
 
+# Function Executors are only listening on localhost so external connections to them are not possible.
+# This is a security measure. Also Executor <-> Function Executor communication is always local and
+# don't support Function Executors running on a different host.
+_FUNCTION_EXECUTOR_SERVER_HOSTNAME = "localhost"
+
+
 def _server_address(port: int) -> str:
-    return f"localhost:{port}"
+    return f"{_FUNCTION_EXECUTOR_SERVER_HOSTNAME}:{port}"
 
 
 def _new_process_group() -> None:
     """Creates a new process group with ID equal to the current process PID. POSIX only."""
     os.setpgid(0, 0)
+
+
+def _find_free_localhost_tcp_port() -> int:
+    with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as sock:
+        sock.bind((_FUNCTION_EXECUTOR_SERVER_HOSTNAME, 0))
+        _, port = sock.getsockname()
+        return port

indexify/executor/function_executor_controller/__init__.py

@@ -0,0 +1,13 @@
+from .function_executor_controller import FunctionExecutorController
+from .loggers import function_executor_logger, task_logger
+from .message_validators import validate_function_executor_description, validate_task
+from .task_output import TaskOutput
+
+__all__ = [
+    "function_executor_logger",
+    "task_logger",
+    "validate_function_executor_description",
+    "validate_task",
+    "FunctionExecutorController",
+    "TaskOutput",
+]

indexify/executor/function_executor_controller/completed_task_metrics.py

@@ -0,0 +1,82 @@
+import time
+from typing import Any
+
+from indexify.proto.executor_api_pb2 import (
+    TaskFailureReason,
+    TaskOutcomeCode,
+)
+
+from .metrics.completed_task_metrics import (
+    METRIC_TASKS_COMPLETED_FAILURE_REASON_ALL,
+    METRIC_TASKS_COMPLETED_FAILURE_REASON_FUNCTION_ERROR,
+    METRIC_TASKS_COMPLETED_FAILURE_REASON_FUNCTION_EXECUTOR_TERMINATED,
+    METRIC_TASKS_COMPLETED_FAILURE_REASON_INTERNAL_ERROR,
+    METRIC_TASKS_COMPLETED_FAILURE_REASON_NONE,
+    METRIC_TASKS_COMPLETED_FAILURE_REASON_TASK_CANCELLED,
+    METRIC_TASKS_COMPLETED_FAILURE_REASON_UNKNOWN,
+    METRIC_TASKS_COMPLETED_OUTCOME_CODE_ALL,
+    METRIC_TASKS_COMPLETED_OUTCOME_CODE_FAILURE,
+    METRIC_TASKS_COMPLETED_OUTCOME_CODE_SUCCESS,
+    metric_task_completion_latency,
+    metric_tasks_completed,
+)
+from .task_info import TaskInfo
+
+
+def emit_completed_task_metrics(task_info: TaskInfo, logger: Any) -> None:
+    """Emits Prometheus metrics for a completed task.
+
+    Doesn't raise any exceptions.
+    """
+    logger = logger.bind(module=__name__)
+    metric_task_completion_latency.observe(time.monotonic() - task_info.start_time)
+
+    task_outcome_code: TaskOutcomeCode = task_info.output.outcome_code
+    task_failure_reason: TaskFailureReason = task_info.output.failure_reason
+    metric_tasks_completed.labels(
+        outcome_code=METRIC_TASKS_COMPLETED_OUTCOME_CODE_ALL,
+        failure_reason=METRIC_TASKS_COMPLETED_FAILURE_REASON_ALL,
+    ).inc()
+    if task_outcome_code == TaskOutcomeCode.TASK_OUTCOME_CODE_SUCCESS:
+        metric_tasks_completed.labels(
+            outcome_code=METRIC_TASKS_COMPLETED_OUTCOME_CODE_SUCCESS,
+            failure_reason=METRIC_TASKS_COMPLETED_FAILURE_REASON_NONE,
+        ).inc()
+    elif task_outcome_code == TaskOutcomeCode.TASK_OUTCOME_CODE_FAILURE:
+        if task_failure_reason == TaskFailureReason.TASK_FAILURE_REASON_INTERNAL_ERROR:
+            metric_tasks_completed.labels(
+                outcome_code=METRIC_TASKS_COMPLETED_OUTCOME_CODE_FAILURE,
+                failure_reason=METRIC_TASKS_COMPLETED_FAILURE_REASON_INTERNAL_ERROR,
+            ).inc()
+        elif (
+            task_failure_reason
+            == TaskFailureReason.TASK_FAILURE_REASON_FUNCTION_EXECUTOR_TERMINATED
+        ):
+            metric_tasks_completed.labels(
+                outcome_code=METRIC_TASKS_COMPLETED_OUTCOME_CODE_FAILURE,
+                failure_reason=METRIC_TASKS_COMPLETED_FAILURE_REASON_FUNCTION_EXECUTOR_TERMINATED,
+            ).inc()
+        elif (
+            task_failure_reason == TaskFailureReason.TASK_FAILURE_REASON_TASK_CANCELLED
+        ):
+            metric_tasks_completed.labels(
+                outcome_code=METRIC_TASKS_COMPLETED_OUTCOME_CODE_FAILURE,
+                failure_reason=METRIC_TASKS_COMPLETED_FAILURE_REASON_TASK_CANCELLED,
+            ).inc()
+        elif task_failure_reason in [
+            TaskFailureReason.TASK_FAILURE_REASON_FUNCTION_ERROR,
+            TaskFailureReason.TASK_FAILURE_REASON_FUNCTION_TIMEOUT,
+        ]:
+            metric_tasks_completed.labels(
+                outcome_code=METRIC_TASKS_COMPLETED_OUTCOME_CODE_FAILURE,
+                failure_reason=METRIC_TASKS_COMPLETED_FAILURE_REASON_FUNCTION_ERROR,
+            ).inc()
+        else:
+            metric_tasks_completed.labels(
+                outcome_code=METRIC_TASKS_COMPLETED_OUTCOME_CODE_FAILURE,
+                failure_reason=METRIC_TASKS_COMPLETED_FAILURE_REASON_UNKNOWN,
+            ).inc()
+            logger.warning(
+                "unexpected task failure reason",
+                failure_reason=TaskFailureReason.Name(task_failure_reason),
+            )

indexify/executor/function_executor_controller/create_function_executor.py

@@ -0,0 +1,158 @@
+import asyncio
+from pathlib import Path
+from typing import Any, Optional
+
+from tensorlake.function_executor.proto.function_executor_pb2 import (
+    InitializeRequest,
+    SerializedObject,
+)
+
+from indexify.executor.blob_store.blob_store import BLOBStore
+from indexify.executor.function_executor.function_executor import (
+    FunctionError,
+    FunctionExecutor,
+    FunctionTimeoutError,
+)
+from indexify.executor.function_executor.server.function_executor_server_factory import (
+    FunctionExecutorServerConfiguration,
+    FunctionExecutorServerFactory,
+)
+from indexify.proto.executor_api_pb2 import (
+    FunctionExecutorDescription,
+    FunctionExecutorTerminationReason,
+)
+
+from .downloads import download_graph
+from .events import FunctionExecutorCreated
+
+
+async def create_function_executor(
+    function_executor_description: FunctionExecutorDescription,
+    function_executor_server_factory: FunctionExecutorServerFactory,
+    blob_store: BLOBStore,
+    executor_id: str,
+    base_url: str,
+    config_path: str,
+    cache_path: Path,
+    logger: Any,
+) -> FunctionExecutorCreated:
+    """Creates a function executor.
+
+    Doesn't raise any exceptions.
+    """
+    logger = logger.bind(module=__name__)
+    try:
+        function_executor: FunctionExecutor = await _create_function_executor(
+            function_executor_description=function_executor_description,
+            function_executor_server_factory=function_executor_server_factory,
+            blob_store=blob_store,
+            executor_id=executor_id,
+            base_url=base_url,
+            config_path=config_path,
+            cache_path=cache_path,
+            logger=logger,
+        )
+        return FunctionExecutorCreated(function_executor)
+    except FunctionTimeoutError as e:
+        return FunctionExecutorCreated(
+            function_executor=None,
+            function_error=e,
+            termination_reason=FunctionExecutorTerminationReason.FUNCTION_EXECUTOR_TERMINATION_REASON_STARTUP_FAILED_FUNCTION_TIMEOUT,
+        )
+    except FunctionError as e:
+        return FunctionExecutorCreated(
+            function_executor=None,
+            function_error=e,
+            termination_reason=FunctionExecutorTerminationReason.FUNCTION_EXECUTOR_TERMINATION_REASON_STARTUP_FAILED_FUNCTION_ERROR,
+        )
+    except BaseException as e:
+        if isinstance(e, asyncio.CancelledError):
+            logger.info("function executor startup was cancelled")
+            return FunctionExecutorCreated(
+                function_executor=None,
+                termination_reason=FunctionExecutorTerminationReason.FUNCTION_EXECUTOR_TERMINATION_REASON_REMOVED_FROM_DESIRED_STATE,
+            )
+        else:
+            logger.error(
+                "failed to create function executor due to platform error",
+                exc_info=e,
+            )
+            return FunctionExecutorCreated(
+                function_executor=None,
+                termination_reason=FunctionExecutorTerminationReason.FUNCTION_EXECUTOR_TERMINATION_REASON_STARTUP_FAILED_INTERNAL_ERROR,
+            )
+
+
+async def _create_function_executor(
+    function_executor_description: FunctionExecutorDescription,
+    function_executor_server_factory: FunctionExecutorServerFactory,
+    blob_store: BLOBStore,
+    executor_id: str,
+    base_url: str,
+    config_path: str,
+    cache_path: Path,
+    logger: Any,
+) -> FunctionExecutor:
+    """Creates a function executor.
+
+    Raises Exception on platform error.
+    Raises FunctionError if customer code failed during FE creation.
+    """
+    graph: SerializedObject = await download_graph(
+        function_executor_description=function_executor_description,
+        cache_path=cache_path,
+        blob_store=blob_store,
+        logger=logger,
+    )
+
+    gpu_count: int = 0
+    if function_executor_description.resources.HasField("gpu"):
+        gpu_count = function_executor_description.resources.gpu.count
+
+    config: FunctionExecutorServerConfiguration = FunctionExecutorServerConfiguration(
+        executor_id=executor_id,
+        function_executor_id=function_executor_description.id,
+        namespace=function_executor_description.namespace,
+        graph_name=function_executor_description.graph_name,
+        graph_version=function_executor_description.graph_version,
+        function_name=function_executor_description.function_name,
+        image_uri=None,
+        secret_names=list(function_executor_description.secret_names),
+        cpu_ms_per_sec=function_executor_description.resources.cpu_ms_per_sec,
+        memory_bytes=function_executor_description.resources.memory_bytes,
+        disk_bytes=function_executor_description.resources.disk_bytes,
+        gpu_count=gpu_count,
+    )
+    if function_executor_description.HasField("image_uri"):
+        config.image_uri = function_executor_description.image_uri
+
+    initialize_request: InitializeRequest = InitializeRequest(
+        namespace=function_executor_description.namespace,
+        graph_name=function_executor_description.graph_name,
+        graph_version=function_executor_description.graph_version,
+        function_name=function_executor_description.function_name,
+        graph=graph,
+    )
+    customer_code_timeout_sec: Optional[float] = None
+    if function_executor_description.HasField("customer_code_timeout_ms"):
+        customer_code_timeout_sec = (
+            function_executor_description.customer_code_timeout_ms / 1000.0
+        )
+
+    function_executor: FunctionExecutor = FunctionExecutor(
+        server_factory=function_executor_server_factory, logger=logger
+    )
+
+    try:
+        # Raises FunctionError if initialization failed in customer code or customer code timed out.
+        await function_executor.initialize(
+            config=config,
+            initialize_request=initialize_request,
+            base_url=base_url,
+            config_path=config_path,
+            customer_code_timeout_sec=customer_code_timeout_sec,
+        )
+        return function_executor
+    except BaseException:  # includes asyncio.CancelledError and anything else
+        await function_executor.destroy()
+        raise

indexify/executor/function_executor_controller/debug_event_loop.py

@@ -0,0 +1,37 @@
+import os
+from typing import Any, List
+
+from .events import BaseEvent
+
+_DEBUG_EVENT_LOOP: bool = (
+    os.getenv("INDEXIFY_FUNCTION_EXECUTOR_CONTROLLER_DEBUG_EVENT_LOOP", "0")
+) == "1"
+
+
+def debug_print_processing_event(event: BaseEvent, logger: Any) -> None:
+    if _DEBUG_EVENT_LOOP:
+        logger.debug(
+            "processing event in control loop",
+            fe_event=str(event),
+        )
+
+
+def debug_print_adding_event(event: BaseEvent, source: str, logger: Any) -> None:
+    if _DEBUG_EVENT_LOOP:
+        logger.debug(
+            "adding event to control loop",
+            source=source,
+            fe_event=str(event),
+        )
+
+
+def debug_print_events(events: List[BaseEvent], logger: Any) -> None:
+    if _DEBUG_EVENT_LOOP:
+        if len(events) == 0:
+            logger.debug("no events n control loop")
+        else:
+            logger.debug(
+                "events in control loop",
+                count=len(events),
+                fe_events=[str(event) for event in events],
+            )

indexify/executor/function_executor_controller/destroy_function_executor.py

@@ -0,0 +1,28 @@
+from typing import Any, Optional
+
+from indexify.executor.function_executor.function_executor import FunctionExecutor
+from indexify.proto.executor_api_pb2 import FunctionExecutorTerminationReason
+
+from .events import FunctionExecutorDestroyed
+
+
+async def destroy_function_executor(
+    function_executor: Optional[FunctionExecutor],
+    termination_reason: FunctionExecutorTerminationReason,
+    logger: Any,
+) -> FunctionExecutorDestroyed:
+    """Destroys a function executor.
+
+    Doesn't raise any exceptions.
+    """
+    logger = logger.bind(module=__name__)
+
+    if function_executor is not None:
+        logger.info(
+            "destroying function executor",
+        )
+        await function_executor.destroy()
+
+    return FunctionExecutorDestroyed(
+        is_success=True, termination_reason=termination_reason
+    )

indexify/executor/function_executor_controller/downloads.py

@@ -0,0 +1,199 @@
+import asyncio
+import os
+from pathlib import Path
+from typing import Any, Optional
+
+import nanoid
+from tensorlake.function_executor.proto.function_executor_pb2 import SerializedObject
+
+from indexify.executor.blob_store.blob_store import BLOBStore
+from indexify.proto.executor_api_pb2 import (
+    DataPayload,
+    DataPayloadEncoding,
+    FunctionExecutorDescription,
+)
+
+from .metrics.downloads import (
+    metric_graph_download_errors,
+    metric_graph_download_latency,
+    metric_graph_downloads,
+    metric_graphs_from_cache,
+    metric_reducer_init_value_download_errors,
+    metric_reducer_init_value_download_latency,
+    metric_reducer_init_value_downloads,
+    metric_task_input_download_errors,
+    metric_task_input_download_latency,
+    metric_task_input_downloads,
+    metric_tasks_downloading_graphs,
+    metric_tasks_downloading_inputs,
+    metric_tasks_downloading_reducer_init_value,
+)
+
+
+async def download_graph(
+    function_executor_description: FunctionExecutorDescription,
+    cache_path: Path,
+    blob_store: BLOBStore,
+    logger: Any,
+) -> SerializedObject:
+    logger = logger.bind(module=__name__)
+    with (
+        metric_graph_download_errors.count_exceptions(),
+        metric_tasks_downloading_graphs.track_inprogress(),
+        metric_graph_download_latency.time(),
+    ):
+        metric_graph_downloads.inc()
+        return await _download_graph(
+            function_executor_description=function_executor_description,
+            cache_path=cache_path,
+            blob_store=blob_store,
+            logger=logger,
+        )
+
+
+async def download_input(
+    data_payload: DataPayload,
+    blob_store: BLOBStore,
+    logger: Any,
+) -> SerializedObject:
+    logger = logger.bind(module=__name__)
+    with (
+        metric_task_input_download_errors.count_exceptions(),
+        metric_tasks_downloading_inputs.track_inprogress(),
+        metric_task_input_download_latency.time(),
+    ):
+        metric_task_input_downloads.inc()
+        return await _download_input(
+            data_payload=data_payload,
+            blob_store=blob_store,
+            logger=logger,
+        )
+
+
+async def download_init_value(
+    data_payload: DataPayload,
+    blob_store: BLOBStore,
+    logger: Any,
+) -> SerializedObject:
+    logger = logger.bind(module=__name__)
+    with (
+        metric_reducer_init_value_download_errors.count_exceptions(),
+        metric_tasks_downloading_reducer_init_value.track_inprogress(),
+        metric_reducer_init_value_download_latency.time(),
+    ):
+        metric_reducer_init_value_downloads.inc()
+        return await _download_input(
+            data_payload=data_payload,
+            blob_store=blob_store,
+            logger=logger,
+        )
+
+
+async def _download_input(
+    data_payload: DataPayload,
+    blob_store: BLOBStore,
+    logger: Any,
+) -> SerializedObject:
+    data: bytes = await blob_store.get(uri=data_payload.uri, logger=logger)
+    return _serialized_object_from_data_payload_proto(
+        data_payload=data_payload,
+        data=data,
+    )
+
+
+async def _download_graph(
+    function_executor_description: FunctionExecutorDescription,
+    cache_path: Path,
+    blob_store: BLOBStore,
+    logger: Any,
+) -> SerializedObject:
+    # Cache graph to reduce load on the server.
+    graph_path = os.path.join(
+        str(cache_path),
+        "graph_cache",
+        function_executor_description.namespace,
+        function_executor_description.graph_name,
+        function_executor_description.graph_version,
+    )
+    # Filesystem operations are synchronous.
+    # Run in a separate thread to not block the main event loop.
+    graph: Optional[SerializedObject] = await asyncio.to_thread(
+        _read_cached_graph, path=graph_path
+    )
+    if graph is not None:
+        metric_graphs_from_cache.inc()
+        return graph
+
+    data: bytes = await blob_store.get(
+        uri=function_executor_description.graph.uri, logger=logger
+    )
+    graph = _serialized_object_from_data_payload_proto(
+        data_payload=function_executor_description.graph,
+        data=data,
+    )
+
+    # Filesystem operations are synchronous.
+    # Run in a separate thread to not block the main event loop.
+    # We don't need to wait for the write completion so we use create_task.
+    asyncio.create_task(
+        asyncio.to_thread(
+            _write_cached_graph, path=graph_path, graph=graph, cache_path=cache_path
+        ),
+        name="graph cache write",
+    )
+
+    return graph
+
+
+def _read_cached_graph(path: str) -> Optional[SerializedObject]:
+    if not os.path.exists(path):
+        return None
+
+    with open(path, "rb") as f:
+        return SerializedObject.FromString(f.read())
+
+
+def _write_cached_graph(path: str, graph: SerializedObject, cache_path: Path) -> None:
+    if os.path.exists(path):
+        # Another task already cached the graph.
+        return None
+
+    tmp_path = os.path.join(str(cache_path), "task_graph_cache", nanoid.generate())
+    os.makedirs(os.path.dirname(tmp_path), exist_ok=True)
+    with open(tmp_path, "wb") as f:
+        f.write(graph.SerializeToString())
+    os.makedirs(os.path.dirname(path), exist_ok=True)
+    # Atomically rename the fully written file at tmp path.
+    # This allows us to not use any locking because file link/unlink
+    # are atomic operations at filesystem level.
+    # This also allows to share the same cache between multiple Executors.
+    os.replace(tmp_path, path)
+
+
+def _serialized_object_from_data_payload_proto(
+    data_payload: DataPayload, data: bytes
+) -> SerializedObject:
+    """Converts the given data payload and its data into SerializedObject accepted by Function Executor.
+
+    Raises ValueError if the supplied data payload can't be converted into serialized object.
+    """
+    if data_payload.encoding == DataPayloadEncoding.DATA_PAYLOAD_ENCODING_BINARY_PICKLE:
+        return SerializedObject(
+            bytes=data,
+            content_type="application/octet-stream",
+        )
+    elif data_payload.encoding == DataPayloadEncoding.DATA_PAYLOAD_ENCODING_UTF8_TEXT:
+        return SerializedObject(
+            content_type="text/plain",
+            string=data.decode("utf-8"),
+        )
+    elif data_payload.encoding == DataPayloadEncoding.DATA_PAYLOAD_ENCODING_UTF8_JSON:
+        result = SerializedObject(
+            content_type="application/json",
+            string=data.decode("utf-8"),
+        )
+        return result
+
+    raise ValueError(
+        f"Can't convert data payload {data_payload} into serialized object"
+    )