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

indexify/executor/executor.py

@@ -12,6 +12,7 @@ from tensorlake.utils.logging import suppress as suppress_logging
 from indexify.proto.executor_api_pb2 import ExecutorStatus
 
 from .api_objects import FunctionURI, Task
+from .blob_store.blob_store import BLOBStore
 from .downloader import Downloader
 from .executor_flavor import ExecutorFlavor
 from .function_executor.function_executor_states_container import (
@@ -64,11 +65,12 @@ class Executor:
         function_allowlist: Optional[List[FunctionURI]],
         function_executor_server_factory: FunctionExecutorServerFactory,
         server_addr: str,
+        grpc_server_addr: str,
         config_path: Optional[str],
         monitoring_server_host: str,
         monitoring_server_port: int,
-        grpc_server_addr: Optional[str],
         enable_grpc_state_reconciler: bool,
+        blob_store: BLOBStore,
     ):
         self._logger = structlog.get_logger(module=__name__)
         self._is_shutdown: bool = False
@@ -95,45 +97,45 @@ class Executor:
             self._function_executor_states
         )
         self._downloader = Downloader(
-            code_path=code_path, base_url=self._base_url, config_path=config_path
+            code_path=code_path,
+            base_url=self._base_url,
+            blob_store=blob_store,
+            config_path=config_path,
+        )
+        self._function_allowlist: Optional[List[FunctionURI]] = function_allowlist
+        self._function_executor_server_factory = function_executor_server_factory
+        self._channel_manager = ChannelManager(
+            server_address=grpc_server_addr,
+            config_path=config_path,
+            logger=self._logger,
+        )
+        self._state_reporter = ExecutorStateReporter(
+            executor_id=id,
+            flavor=flavor,
+            version=version,
+            labels=labels,
+            development_mode=development_mode,
+            function_allowlist=self._function_allowlist,
+            function_executor_states=self._function_executor_states,
+            channel_manager=self._channel_manager,
+            logger=self._logger,
+        )
+        self._state_reporter.update_executor_status(
+            ExecutorStatus.EXECUTOR_STATUS_STARTING_UP
         )
         self._task_reporter = TaskReporter(
             base_url=self._base_url,
             executor_id=id,
             config_path=config_path,
+            channel_manager=self._channel_manager,
         )
-        self._function_allowlist: Optional[List[FunctionURI]] = function_allowlist
-        self._function_executor_server_factory = function_executor_server_factory
 
-        # HTTP mode services
+        # HTTP mode task runner
         self._task_runner: Optional[TaskRunner] = None
         self._task_fetcher: Optional[TaskFetcher] = None
-        # gRPC mode services
-        self._channel_manager: Optional[ChannelManager] = None
-        self._state_reporter: Optional[ExecutorStateReporter] = None
+        # gRPC mode state reconciler that runs tasks
         self._state_reconciler: Optional[ExecutorStateReconciler] = None
 
-        if grpc_server_addr is not None:
-            self._channel_manager = ChannelManager(
-                server_address=grpc_server_addr,
-                config_path=config_path,
-                logger=self._logger,
-            )
-            self._state_reporter = ExecutorStateReporter(
-                executor_id=id,
-                flavor=flavor,
-                version=version,
-                labels=labels,
-                development_mode=development_mode,
-                function_allowlist=self._function_allowlist,
-                function_executor_states=self._function_executor_states,
-                channel_manager=self._channel_manager,
-                logger=self._logger,
-            )
-            self._state_reporter.update_executor_status(
-                ExecutorStatus.EXECUTOR_STATUS_STARTING_UP
-            )
-
         if enable_grpc_state_reconciler:
             self._state_reconciler = ExecutorStateReconciler(
                 executor_id=id,
@@ -171,8 +173,8 @@ class Executor:
             "version": version,
             "code_path": str(code_path),
             "server_addr": server_addr,
-            "config_path": str(config_path),
             "grpc_server_addr": str(grpc_server_addr),
+            "config_path": str(config_path),
             "enable_grpc_state_reconciler": str(enable_grpc_state_reconciler),
             "hostname": gethostname(),
         }
@@ -192,12 +194,15 @@ class Executor:
                 signum, self.shutdown, asyncio.get_event_loop()
             )
 
-        asyncio.get_event_loop().create_task(self._monitoring_server.run())
-        if self._state_reporter is not None:
-            self._state_reporter.update_executor_status(
-                ExecutorStatus.EXECUTOR_STATUS_RUNNING
-            )
-            asyncio.get_event_loop().create_task(self._state_reporter.run())
+        asyncio.get_event_loop().create_task(
+            self._monitoring_server.run(), name="monitoring server runner"
+        )
+        self._state_reporter.update_executor_status(
+            ExecutorStatus.EXECUTOR_STATUS_RUNNING
+        )
+        asyncio.get_event_loop().create_task(
+            self._state_reporter.run(), name="state reporter runner"
+        )
 
         metric_executor_state.state("running")
         self._startup_probe_handler.set_ready()
@@ -218,7 +223,6 @@ class Executor:
         """Runs the gRPC state reconciler and state reporter.
 
         Never raises any exceptions."""
-        asyncio.create_task(self._state_reporter.run())
         await self._state_reconciler.run()
 
     async def _http_task_runner_loop(self):
@@ -227,11 +231,15 @@ class Executor:
                 async for task in self._task_fetcher.run():
                     metric_tasks_fetched.inc()
                     if not self._is_shutdown:
-                        asyncio.create_task(self._run_task(task))
+                        asyncio.create_task(
+                            self._run_task(task), name="task runner (http mode)"
+                        )
+                self._logger.info("fetching tasks finished, reconnecting in 5 seconds")
             except Exception as e:
                 self._logger.error(
                     "failed fetching tasks, retrying in 5 seconds", exc_info=e
                 )
+            if not self._is_shutdown:
                 await asyncio.sleep(5)
 
     async def _run_task(self, task: Task) -> None:
@@ -256,6 +264,9 @@ class Executor:
             )
             logger.error("task execution failed", exc_info=e)
 
+        if output.metrics is not None:
+            self.log_function_metrics(output)
+
         with (
             metric_tasks_reporting_outcome.track_inprogress(),
             metric_task_outcome_report_latency.time(),
@@ -265,18 +276,42 @@ class Executor:
 
         metric_task_completion_latency.observe(time.monotonic() - start_time)
 
+    def log_function_metrics(self, output: TaskOutput):
+        for counter_name, counter_value in output.metrics.counters.items():
+            self._logger.info(
+                f"function_metric",
+                counter_name=counter_name,
+                counter_value=counter_value,
+                invocation_id=output.graph_invocation_id,
+                function_name=output.function_name,
+                graph_name=output.graph_name,
+                namespace=output.namespace,
+            )
+        for timer_name, timer_value in output.metrics.timers.items():
+            self._logger.info(
+                f"function_metric",
+                timer_name=timer_name,
+                timer_value=timer_value,
+                invocation_id=output.graph_invocation_id,
+                function_name=output.function_name,
+                graph_name=output.graph_name,
+                namespace=output.namespace,
+            )
+
     async def _run_task_and_get_output(self, task: Task, logger: Any) -> TaskOutput:
         graph: SerializedObject = await self._downloader.download_graph(
             namespace=task.namespace,
             graph_name=task.compute_graph,
             graph_version=task.graph_version,
             logger=logger,
+            data_payload=None,
         )
         input: SerializedObject = await self._downloader.download_input(
             namespace=task.namespace,
             graph_name=task.compute_graph,
             graph_invocation_id=task.invocation_id,
             input_key=task.input_key,
+            data_payload=None,
             logger=logger,
         )
         init_value: Optional[SerializedObject] = (
@@ -289,6 +324,7 @@ class Executor:
                     function_name=task.compute_fn,
                     graph_invocation_id=task.invocation_id,
                     reducer_output_key=task.reducer_output_id,
+                    data_payload=None,
                     logger=logger,
                 )
             )
@@ -358,12 +394,12 @@ class Executor:
         if self._task_runner is not None:
             await self._task_runner.shutdown()
 
-        if self._channel_manager is not None:
-            await self._channel_manager.shutdown()
         if self._state_reporter is not None:
             await self._state_reporter.shutdown()
         if self._state_reconciler is not None:
             await self._state_reconciler.shutdown()
+        if self._channel_manager is not None:
+            await self._channel_manager.destroy()
 
         # We need to shutdown all users of FE states first,
         # otherwise states might disappear unexpectedly and we might
@@ -375,7 +411,7 @@ class Executor:
         # The current task is cancelled, the code after this line will not run.
 
     def shutdown(self, loop):
-        loop.create_task(self._shutdown(loop))
+        loop.create_task(self._shutdown(loop), name="executor shutdown")
 
     def _task_logger(self, task: Task) -> Any:
         return self._logger.bind(

indexify/executor/function_executor/function_executor.py

@@ -88,6 +88,7 @@ class FunctionExecutor:
         initialize_request: InitializeRequest,
         base_url: str,
         config_path: Optional[str],
+        customer_code_timeout_sec: Optional[float] = None,
     ):
         """Creates and initializes a FunctionExecutorServer and all resources associated with it.
 
@@ -103,7 +104,9 @@ class FunctionExecutor:
                 await self._establish_channel()
                 stub: FunctionExecutorStub = FunctionExecutorStub(self._channel)
                 await _collect_server_info(stub)
-                await _initialize_server(stub, initialize_request)
+                await _initialize_server(
+                    stub, initialize_request, customer_code_timeout_sec
+                )
                 await self._create_invocation_state_client(
                     stub=stub,
                     base_url=base_url,
@@ -293,18 +296,28 @@ async def _collect_server_info(stub: FunctionExecutorStub) -> None:
 
 
 async def _initialize_server(
-    stub: FunctionExecutorStub, initialize_request: InitializeRequest
+    stub: FunctionExecutorStub,
+    initialize_request: InitializeRequest,
+    customer_code_timeout_sec: Optional[float],
 ) -> None:
     with (
         metric_initialize_rpc_errors.count_exceptions(),
         metric_initialize_rpc_latency.time(),
     ):
-        initialize_response: InitializeResponse = await stub.initialize(
-            initialize_request
-        )
-        if initialize_response.success:
-            return
-        if initialize_response.HasField("customer_error"):
-            raise CustomerError(initialize_response.customer_error)
-        else:
-            raise Exception("initialize RPC failed at function executor server")
+        try:
+            initialize_response: InitializeResponse = await stub.initialize(
+                initialize_request,
+                timeout=customer_code_timeout_sec,
+            )
+            if initialize_response.success:
+                return
+            if initialize_response.HasField("customer_error"):
+                raise CustomerError(initialize_response.customer_error)
+            else:
+                raise Exception("initialize RPC failed at function executor server")
+        except grpc.aio.AioRpcError as e:
+            if e.code() == grpc.StatusCode.DEADLINE_EXCEEDED:
+                raise CustomerError(
+                    f"Customer code timeout of {customer_code_timeout_sec:.3f} sec expired"
+                ) from e
+            raise

indexify/executor/function_executor/function_executor_state.py

@@ -25,6 +25,7 @@ class FunctionExecutorState:
         graph_version: str,
         function_name: str,
         image_uri: Optional[str],
+        secret_names: List[str],
         logger: Any,
     ):
         # Read only fields.
@@ -33,6 +34,7 @@ class FunctionExecutorState:
         self.graph_name: str = graph_name
         self.function_name: str = function_name
         self.image_uri: Optional[str] = image_uri
+        self.secret_names: List[str] = secret_names
         self._logger: Any = logger.bind(
             module=__name__,
             function_executor_id=id,
@@ -47,6 +49,7 @@ class FunctionExecutorState:
         # TODO: Move graph_version to immutable fields once we migrate to gRPC State Reconciler.
         self.graph_version: str = graph_version
         self.status: FunctionExecutorStatus = FunctionExecutorStatus.DESTROYED
+        self.status_message: str = ""
         self.status_change_notifier: asyncio.Condition = asyncio.Condition(
             lock=self.lock
         )
@@ -62,7 +65,9 @@ class FunctionExecutorState:
         while self.status not in allowlist:
             await self.status_change_notifier.wait()
 
-    async def set_status(self, new_status: FunctionExecutorStatus) -> None:
+    async def set_status(
+        self, new_status: FunctionExecutorStatus, status_message: str = ""
+    ) -> None:
         """Sets the status of the Function Executor.
 
         The caller must hold the lock.
@@ -70,6 +75,7 @@ class FunctionExecutorState:
         """
         self.check_locked()
         if is_status_change_allowed(self.status, new_status):
+            # If status didn't change then still log it for visibility.
             self._logger.info(
                 "function executor status changed",
                 old_status=self.status.name,
@@ -78,12 +84,14 @@ class FunctionExecutorState:
             metric_function_executors_with_status.labels(status=self.status.name).dec()
             metric_function_executors_with_status.labels(status=new_status.name).inc()
             self.status = new_status
+            self.status_message = status_message
             self.status_change_notifier.notify_all()
         else:
             raise ValueError(
                 f"Invalid status change from {self.status} to {new_status}"
             )
 
+    # TODO: Delete this method once HTTP protocol is removed as it's used only there.
     async def destroy_function_executor(self) -> None:
         """Destroys the Function Executor if it exists.
 

indexify/executor/function_executor/function_executor_states_container.py

@@ -1,5 +1,5 @@
 import asyncio
-from typing import Any, AsyncGenerator, Dict, Optional
+from typing import Any, AsyncGenerator, Dict, List, Optional
 
 from .function_executor_state import FunctionExecutorState
 from .function_executor_status import FunctionExecutorStatus
@@ -26,6 +26,7 @@ class FunctionExecutorStatesContainer:
         graph_version: str,
         function_name: str,
         image_uri: Optional[str],
+        secret_names: List[str],
     ) -> FunctionExecutorState:
         """Get or create a function executor state with the given ID.
 
@@ -45,6 +46,7 @@ class FunctionExecutorStatesContainer:
                     graph_version=graph_version,
                     function_name=function_name,
                     image_uri=image_uri,
+                    secret_names=secret_names,
                     logger=self._logger,
                 )
                 self._states[id] = state
@@ -52,6 +54,11 @@ class FunctionExecutorStatesContainer:
 
             return self._states[id]
 
+    async def get(self, id: str) -> FunctionExecutorState:
+        """Get the state with the given ID. Raises Exception if the state does not exist."""
+        async with self._lock:
+            return self._states[id]
+
     async def __aiter__(self) -> AsyncGenerator[FunctionExecutorState, None]:
         async with self._lock:
             for state in self._states.values():

indexify/executor/function_executor/function_executor_status.py

@@ -23,6 +23,7 @@ class FunctionExecutorStatus(Enum):
     UNHEALTHY = "Unhealthy"
     # STARTUP_FAILED_CUSTOMER_ERROR -> DESTROYING
     # STARTUP_FAILED_PLATFORM_ERROR -> DESTROYING
+    # RUNNING_TASK -> DESTROYING
     # UNHEALTHY -> DESTROYING
     # IDLE -> DESTROYING
     DESTROYING = "Destroying"
@@ -33,6 +34,8 @@ class FunctionExecutorStatus(Enum):
     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:
@@ -69,6 +72,7 @@ def is_status_change_allowed(
         ],
         FunctionExecutorStatus.RUNNING_TASK: [
             FunctionExecutorStatus.RUNNING_TASK,
+            FunctionExecutorStatus.DESTROYING,
             FunctionExecutorStatus.IDLE,
             FunctionExecutorStatus.UNHEALTHY,
             FunctionExecutorStatus.SHUTDOWN,

indexify/executor/function_executor/health_checker.py

@@ -107,7 +107,9 @@ class HealthChecker:
             return
 
         self._health_check_failed_callback = callback
-        self._health_check_loop_task = asyncio.create_task(self._health_check_loop())
+        self._health_check_loop_task = asyncio.create_task(
+            self._health_check_loop(), name="function executor health checker loop"
+        )
 
     def stop(self) -> None:
         """Stops the periodic health checks.
@@ -126,7 +128,10 @@ class HealthChecker:
                 break
             await asyncio.sleep(HEALTH_CHECK_POLL_PERIOD_SEC)
 
-        asyncio.create_task(self._health_check_failed_callback(result))
+        asyncio.create_task(
+            self._health_check_failed_callback(result),
+            name="function executor health check failure callback",
+        )
         self._health_check_loop_task = None
 
 

indexify/executor/function_executor/invocation_state_client.py

@@ -67,7 +67,8 @@ class InvocationStateClient:
             self._response_generator()
         )
         self._request_loop_task = asyncio.create_task(
-            self._request_loop(server_requests)
+            self._request_loop(server_requests),
+            name="graph invocation state client request processing loop",
         )
 
     def add_task_to_invocation_id_entry(self, task_id: str, invocation_id: str) -> None:
@@ -100,7 +101,8 @@ class InvocationStateClient:
             pass
         except asyncio.CancelledError:
             # This async task was cancelled by destroy(). Normal situation too.
-            pass
+            # This exception should not be suppressed, see Python asyncio docs.
+            raise
         except Exception as e:
             metric_request_read_errors.inc()
             self._logger.error(

indexify/executor/function_executor/server/subprocess_function_executor_server_factory.py

@@ -25,6 +25,12 @@ class SubprocessFunctionExecutorServerFactory(FunctionExecutorServerFactory):
         logger = logger.bind(module=__name__)
         port: Optional[int] = None
 
+        if len(config.secret_names) > 0:
+            logger.warning(
+                "Subprocess Function Executor does not support secrets. Please supply secrets as environment variables.",
+                secret_names=config.secret_names,
+            )
+
         try:
             port = self._allocate_port()
             args = [

indexify/executor/function_executor/single_task_runner.py

@@ -10,6 +10,7 @@ from tensorlake.function_executor.proto.function_executor_pb2 import (
 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
@@ -26,7 +27,7 @@ from .server.function_executor_server_factory import (
     FunctionExecutorServerFactory,
 )
 from .task_input import TaskInput
-from .task_output import TaskOutput
+from .task_output import TaskMetrics, TaskOutput
 
 
 class SingleTaskRunner:
@@ -286,16 +287,17 @@ class _RunningTaskContextManager:
 
 
 def _task_output(task: Task, response: RunTaskResponse) -> TaskOutput:
-    required_fields = [
-        "stdout",
-        "stderr",
-        "is_reducer",
-        "success",
-    ]
-
-    for field in required_fields:
-        if not response.HasField(field):
-            raise ValueError(f"Response is missing required field: {field}")
+    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,
@@ -308,10 +310,12 @@ def _task_output(task: Task, response: RunTaskResponse) -> TaskOutput:
         stderr=response.stderr,
         reducer=response.is_reducer,
         success=response.success,
+        metrics=metrics,
     )
 
     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
 

indexify/executor/function_executor/task_output.py

@@ -1,11 +1,17 @@
-from typing import Optional
+from typing import Dict, Optional
 
 from tensorlake.function_executor.proto.function_executor_pb2 import (
     FunctionOutput,
     RouterOutput,
 )
 
-from ..api_objects import Task
+
+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:
@@ -19,6 +25,7 @@ class TaskOutput:
         function_name: str,
         graph_version: str,
         graph_invocation_id: str,
+        output_encoding: Optional[str] = None,
         function_output: Optional[FunctionOutput] = None,
         router_output: Optional[RouterOutput] = None,
         stdout: Optional[str] = None,
@@ -26,6 +33,7 @@ class TaskOutput:
         reducer: bool = False,
         success: bool = False,
         is_internal_error: bool = False,
+        metrics: Optional[TaskMetrics] = None,
     ):
         self.task_id = task_id
         self.namespace = namespace
@@ -40,6 +48,8 @@ class TaskOutput:
         self.reducer = reducer
         self.success = success
         self.is_internal_error = is_internal_error
+        self.metrics = metrics
+        self.output_encoding = output_encoding
 
     @classmethod
     def internal_error(
@@ -63,3 +73,27 @@ class TaskOutput:
             stderr="Platform failed to execute the function.",
             is_internal_error=True,
         )
+
+    @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,
+    ) -> "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 exceeded its configured timeout of {timeout_sec:.3f} sec.",
+            is_internal_error=False,
+        )

indexify/executor/grpc/channel_manager.py

@@ -69,6 +69,10 @@ class ChannelManager:
             certificate_chain=certificate_chain,
         )
 
+    async def destroy(self):
+        if self._channel is not None:
+            await self._destroy_locked_channel()
+
     async def get_channel(self) -> grpc.aio.Channel:
         """Returns a channel to the gRPC server.
 
@@ -155,6 +159,3 @@ class ChannelManager:
         except Exception as e:
             self._logger.error("failed closing channel", exc_info=e)
         self._channel = None
-
-    async def shutdown(self):
-        pass