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

indexify/executor/grpc/state_reconciler.py

@@ -7,14 +7,14 @@ from tensorlake.function_executor.proto.function_executor_pb2 import (
     SerializedObject,
 )
 
-from indexify.proto.task_scheduler_pb2 import (
+from indexify.proto.executor_api_pb2 import (
     DesiredExecutorState,
     FunctionExecutorDescription,
     FunctionExecutorStatus,
     GetDesiredExecutorStatesRequest,
 )
-from indexify.proto.task_scheduler_pb2_grpc import (
-    TaskSchedulerServiceStub,
+from indexify.proto.executor_api_pb2_grpc import (
+    ExecutorAPIStub,
 )
 
 from ..downloader import Downloader
@@ -30,20 +30,11 @@ from ..function_executor.server.function_executor_server_factory import (
 from ..function_executor.task_input import TaskInput
 from ..function_executor.task_output import TaskOutput
 from ..metrics.executor import (
-    METRIC_TASKS_COMPLETED_OUTCOME_ALL,
-    METRIC_TASKS_COMPLETED_OUTCOME_ERROR_CUSTOMER_CODE,
-    METRIC_TASKS_COMPLETED_OUTCOME_ERROR_PLATFORM,
-    METRIC_TASKS_COMPLETED_OUTCOME_SUCCESS,
-    metric_task_completion_latency,
-    metric_task_outcome_report_latency,
-    metric_task_outcome_report_retries,
-    metric_task_outcome_reports,
-    metric_tasks_completed,
     metric_tasks_fetched,
-    metric_tasks_reporting_outcome,
 )
 from ..task_reporter import TaskReporter
-from .channel_creator import ChannelCreator
+from .channel_manager import ChannelManager
+from .state_reporter import ExecutorStateReporter
 
 _RECONCILE_STREAM_BACKOFF_INTERVAL_SEC = 5
 
@@ -58,7 +49,8 @@ class ExecutorStateReconciler:
         config_path: Optional[str],
         downloader: Downloader,
         task_reporter: TaskReporter,
-        channel_creator: ChannelCreator,
+        channel_manager: ChannelManager,
+        state_reporter: ExecutorStateReporter,
         logger: Any,
     ):
         self._executor_id: str = executor_id
@@ -72,7 +64,8 @@ class ExecutorStateReconciler:
         self._function_executor_states: FunctionExecutorStatesContainer = (
             function_executor_states
         )
-        self._channel_creator = channel_creator
+        self._channel_manager: ChannelManager = channel_manager
+        self._state_reporter: ExecutorStateReporter = state_reporter
         self._logger: Any = logger.bind(module=__name__)
         self._is_shutdown: bool = False
         self._server_last_clock: Optional[int] = None
@@ -83,27 +76,25 @@ class ExecutorStateReconciler:
         Never raises any exceptions.
         """
         while not self._is_shutdown:
-            async with await self._channel_creator.create() as server_channel:
-                server_channel: grpc.aio.Channel
-                stub = TaskSchedulerServiceStub(server_channel)
-                while not self._is_shutdown:
-                    try:
-                        # TODO: Report state once before starting the stream.
-                        desired_states_stream: AsyncGenerator[
-                            DesiredExecutorState, None
-                        ] = stub.get_desired_executor_states(
-                            GetDesiredExecutorStatesRequest(
-                                executor_id=self._executor_id
-                            )
-                        )
-                        await self._process_desired_states_stream(desired_states_stream)
-                    except Exception as e:
-                        self._logger.error(
-                            f"Failed processing desired states stream, reconnecting in {_RECONCILE_STREAM_BACKOFF_INTERVAL_SEC} sec.",
-                            exc_info=e,
-                        )
-                        await asyncio.sleep(_RECONCILE_STREAM_BACKOFF_INTERVAL_SEC)
-                        break
+            stub = ExecutorAPIStub(await self._channel_manager.get_channel())
+            while not self._is_shutdown:
+                try:
+                    # Report state once before starting the stream so Server
+                    # doesn't use old state it knew about this Executor in the past.
+                    await self._state_reporter.report_state(stub)
+                    desired_states_stream: AsyncGenerator[
+                        DesiredExecutorState, None
+                    ] = stub.get_desired_executor_states(
+                        GetDesiredExecutorStatesRequest(executor_id=self._executor_id)
+                    )
+                    await self._process_desired_states_stream(desired_states_stream)
+                except Exception as e:
+                    self._logger.error(
+                        f"Failed processing desired states stream, reconnecting in {_RECONCILE_STREAM_BACKOFF_INTERVAL_SEC} sec.",
+                        exc_info=e,
+                    )
+                    await asyncio.sleep(_RECONCILE_STREAM_BACKOFF_INTERVAL_SEC)
+                    break
 
         self._logger.info("State reconciler shutdown.")
 
@@ -123,6 +114,7 @@ class ExecutorStateReconciler:
             await self._reconcile_state(new_state)
 
     async def _reconcile_state(self, new_state: DesiredExecutorState):
+        # TODO: use completed_tasks_container to ignore tasks that were already completed.
         await self._reconcile_function_executors(new_state)
         # TODO
         # await self._reconcile_task_allocations(new_state)
@@ -148,6 +140,7 @@ class ExecutorStateReconciler:
                     graph_version=desired_function_executor.graph_version,
                     function_name=desired_function_executor.function_name,
                     image_uri=desired_function_executor.image_uri,
+                    secret_names=list(desired_function_executor.secret_names),
                 )
             )
 
@@ -291,7 +284,9 @@ class ExecutorStateReconciler:
         while True:
             logger = logger.bind(retries=reporting_retries)
             try:
-                await self._task_reporter.report(output=task_output, logger=logger)
+                await self._task_reporter.report(
+                    data_payload=task_output, logger=logger
+                )
                 break
             except Exception as e:
                 logger.error(

indexify/executor/grpc/state_reporter.py

@@ -1,37 +1,44 @@
 import asyncio
+import hashlib
+from socket import gethostname
 from typing import Any, Dict, List, Optional
 
 import grpc
 
-from indexify.proto.task_scheduler_pb2 import (
+from indexify.proto.executor_api_pb2 import (
     AllowedFunction,
+)
+from indexify.proto.executor_api_pb2 import ExecutorFlavor as ExecutorFlavorProto
+from indexify.proto.executor_api_pb2 import (
     ExecutorState,
     ExecutorStatus,
     FunctionExecutorDescription,
 )
-from indexify.proto.task_scheduler_pb2 import (
+from indexify.proto.executor_api_pb2 import (
     FunctionExecutorState as FunctionExecutorStateProto,
 )
-from indexify.proto.task_scheduler_pb2 import (
+from indexify.proto.executor_api_pb2 import (
     FunctionExecutorStatus as FunctionExecutorStatusProto,
 )
-from indexify.proto.task_scheduler_pb2 import (
+from indexify.proto.executor_api_pb2 import (
     GPUModel,
     GPUResources,
     HostResources,
     ReportExecutorStateRequest,
 )
-from indexify.proto.task_scheduler_pb2_grpc import (
-    TaskSchedulerServiceStub,
+from indexify.proto.executor_api_pb2_grpc import (
+    ExecutorAPIStub,
 )
 
 from ..api_objects import FunctionURI
+from ..executor_flavor import ExecutorFlavor
 from ..function_executor.function_executor_state import FunctionExecutorState
 from ..function_executor.function_executor_states_container import (
     FunctionExecutorStatesContainer,
 )
 from ..function_executor.function_executor_status import FunctionExecutorStatus
-from .channel_creator import ChannelCreator
+from ..runtime_probes import RuntimeProbes
+from .channel_manager import ChannelManager
 from .metrics.state_reporter import (
     metric_state_report_errors,
     metric_state_report_latency,
@@ -47,24 +54,32 @@ class ExecutorStateReporter:
     def __init__(
         self,
         executor_id: str,
+        flavor: ExecutorFlavor,
+        version: str,
+        labels: Dict[str, str],
         development_mode: bool,
         function_allowlist: Optional[List[FunctionURI]],
         function_executor_states: FunctionExecutorStatesContainer,
-        channel_creator: ChannelCreator,
+        channel_manager: ChannelManager,
         logger: Any,
     ):
         self._executor_id: str = executor_id
+        self._flavor: ExecutorFlavor = flavor
+        self._version: str = version
+        self._labels: Dict[str, str] = labels.copy()
         self._development_mode: bool = development_mode
+        self._hostname: str = gethostname()
         self._function_executor_states: FunctionExecutorStatesContainer = (
             function_executor_states
         )
-        self._channel_creator = channel_creator
+        self._channel_manager = channel_manager
         self._logger: Any = logger.bind(module=__name__)
         self._is_shutdown: bool = False
         self._executor_status: ExecutorStatus = ExecutorStatus.EXECUTOR_STATUS_UNKNOWN
         self._allowed_functions: List[AllowedFunction] = _to_grpc_allowed_functions(
             function_allowlist
         )
+        self._labels.update(_label_values_to_strings(RuntimeProbes().probe().labels))
 
     def update_executor_status(self, value: ExecutorStatus):
         self._executor_status = value
@@ -75,24 +90,30 @@ class ExecutorStateReporter:
         Never raises any exceptions.
         """
         while not self._is_shutdown:
-            async with await self._channel_creator.create() as server_channel:
-                server_channel: grpc.aio.Channel
-                stub = TaskSchedulerServiceStub(server_channel)
-                while not self._is_shutdown:
-                    try:
-                        await self._report_state(stub)
-                        await asyncio.sleep(_REPORTING_INTERVAL_SEC)
-                    except Exception as e:
-                        self._logger.error(
-                            f"Failed to report state to the server, reconnecting in {_REPORT_BACKOFF_ON_ERROR_SEC} sec.",
-                            exc_info=e,
-                        )
-                        await asyncio.sleep(_REPORT_BACKOFF_ON_ERROR_SEC)
-                        break
+            stub = ExecutorAPIStub(await self._channel_manager.get_channel())
+            while not self._is_shutdown:
+                try:
+                    # The periodic state reports serve as channel health monitoring requests
+                    # (same as TCP keep-alive). Channel Manager returns the same healthy channel
+                    # for all RPCs that we do from Executor to Server. So all the RPCs benefit
+                    # from this channel health monitoring.
+                    await self.report_state(stub)
+                    await asyncio.sleep(_REPORTING_INTERVAL_SEC)
+                except Exception as e:
+                    self._logger.error(
+                        f"Failed to report state to the server, reconnecting in {_REPORT_BACKOFF_ON_ERROR_SEC} sec.",
+                        exc_info=e,
+                    )
+                    await asyncio.sleep(_REPORT_BACKOFF_ON_ERROR_SEC)
+                    break
 
         self._logger.info("State reporter shutdown")
 
-    async def _report_state(self, stub: TaskSchedulerServiceStub):
+    async def report_state(self, stub: ExecutorAPIStub):
+        """Reports the current state to the server represented by the supplied stub.
+
+        Raises exceptions on failure.
+        """
         with (
             metric_state_report_errors.count_exceptions(),
             metric_state_report_latency.time(),
@@ -101,11 +122,16 @@ class ExecutorStateReporter:
             state = ExecutorState(
                 executor_id=self._executor_id,
                 development_mode=self._development_mode,
-                executor_status=self._executor_status,
+                hostname=self._hostname,
+                flavor=_to_grpc_executor_flavor(self._flavor, self._logger),
+                version=self._version,
+                status=self._executor_status,
                 free_resources=await self._fetch_free_host_resources(),
                 allowed_functions=self._allowed_functions,
                 function_executor_states=await self._fetch_function_executor_states(),
+                labels=self._labels,
             )
+            state.state_hash = _state_hash(state)
 
             await stub.report_executor_state(
                 ReportExecutorStateRequest(executor_state=state),
@@ -129,20 +155,25 @@ class ExecutorStateReporter:
 
         async for function_executor_state in self._function_executor_states:
             function_executor_state: FunctionExecutorState
-            states.append(
-                FunctionExecutorStateProto(
-                    description=FunctionExecutorDescription(
-                        id=function_executor_state.id,
-                        namespace=function_executor_state.namespace,
-                        graph_name=function_executor_state.graph_name,
-                        graph_version=function_executor_state.graph_version,
-                        function_name=function_executor_state.function_name,
-                    ),
-                    status=_to_grpc_function_executor_status(
-                        function_executor_state.status, self._logger
-                    ),
-                )
+            function_executor_state_proto = FunctionExecutorStateProto(
+                description=FunctionExecutorDescription(
+                    id=function_executor_state.id,
+                    namespace=function_executor_state.namespace,
+                    graph_name=function_executor_state.graph_name,
+                    graph_version=function_executor_state.graph_version,
+                    function_name=function_executor_state.function_name,
+                    secret_names=function_executor_state.secret_names,
+                ),
+                status=_to_grpc_function_executor_status(
+                    function_executor_state.status, self._logger
+                ),
+                status_message=function_executor_state.status_message,
             )
+            if function_executor_state.image_uri:
+                function_executor_state_proto.description.image_uri = (
+                    function_executor_state.image_uri
+                )
+            states.append(function_executor_state_proto)
 
         return states
 
@@ -182,7 +213,7 @@ _STATUS_MAPPING: Dict[FunctionExecutorStatus, Any] = {
     FunctionExecutorStatus.UNHEALTHY: FunctionExecutorStatusProto.FUNCTION_EXECUTOR_STATUS_UNHEALTHY,
     FunctionExecutorStatus.DESTROYING: FunctionExecutorStatusProto.FUNCTION_EXECUTOR_STATUS_STOPPING,
     FunctionExecutorStatus.DESTROYED: FunctionExecutorStatusProto.FUNCTION_EXECUTOR_STATUS_STOPPED,
-    FunctionExecutorStatus.SHUTDOWN: FunctionExecutorStatusProto.FUNCTION_EXECUTOR_STATUS_STOPPED,
+    FunctionExecutorStatus.SHUTDOWN: FunctionExecutorStatusProto.FUNCTION_EXECUTOR_STATUS_SHUTDOWN,
 }
 
 
@@ -197,3 +228,33 @@ def _to_grpc_function_executor_status(
         logger.error("Unexpected Function Executor status", status=status)
 
     return result
+
+
+_FLAVOR_MAPPING = {
+    ExecutorFlavor.OSS: ExecutorFlavorProto.EXECUTOR_FLAVOR_OSS,
+    ExecutorFlavor.PLATFORM: ExecutorFlavorProto.EXECUTOR_FLAVOR_PLATFORM,
+}
+
+
+def _to_grpc_executor_flavor(
+    flavor: ExecutorFlavor, logger: Any
+) -> ExecutorFlavorProto:
+    result: ExecutorFlavorProto = _FLAVOR_MAPPING.get(
+        flavor, ExecutorFlavorProto.EXECUTOR_FLAVOR_UNKNOWN
+    )
+
+    if result == ExecutorFlavorProto.EXECUTOR_FLAVOR_UNKNOWN:
+        logger.error("Unexpected Executor flavor", flavor=flavor)
+
+    return result
+
+
+def _label_values_to_strings(labels: Dict[str, Any]) -> Dict[str, str]:
+    return {k: str(v) for k, v in labels.items()}
+
+
+def _state_hash(state: ExecutorState) -> str:
+    serialized_state: bytes = state.SerializeToString(deterministic=True)
+    hasher = hashlib.sha256(usedforsecurity=False)
+    hasher.update(serialized_state)
+    return hasher.hexdigest()