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

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

@@ -7,29 +7,29 @@ from tensorlake.function_executor.proto.function_executor_pb2 import (
     SerializedObject,
 )
 
-from indexify.task_scheduler.proto.task_scheduler_pb2 import (
+from indexify.proto.task_scheduler_pb2 import (
     DesiredExecutorState,
     FunctionExecutorDescription,
     FunctionExecutorStatus,
     GetDesiredExecutorStatesRequest,
 )
-from indexify.task_scheduler.proto.task_scheduler_pb2_grpc import (
+from indexify.proto.task_scheduler_pb2_grpc import (
     TaskSchedulerServiceStub,
 )
 
-from .downloader import Downloader
-from .function_executor.function_executor import CustomerError, FunctionExecutor
-from .function_executor.function_executor_state import FunctionExecutorState
-from .function_executor.function_executor_states_container import (
+from ..downloader import Downloader
+from ..function_executor.function_executor import CustomerError, FunctionExecutor
+from ..function_executor.function_executor_state import FunctionExecutorState
+from ..function_executor.function_executor_states_container import (
     FunctionExecutorStatesContainer,
 )
-from .function_executor.server.function_executor_server_factory import (
+from ..function_executor.server.function_executor_server_factory import (
     FunctionExecutorServerConfiguration,
     FunctionExecutorServerFactory,
 )
-from .function_executor.task_input import TaskInput
-from .function_executor.task_output import TaskOutput
-from .metrics.executor 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,
@@ -42,7 +42,10 @@ from .metrics.executor import (
     metric_tasks_fetched,
     metric_tasks_reporting_outcome,
 )
-from .task_reporter import TaskReporter
+from ..task_reporter import TaskReporter
+from .channel_creator import ChannelCreator
+
+_RECONCILE_STREAM_BACKOFF_INTERVAL_SEC = 5
 
 
 class ExecutorStateReconciler:
@@ -55,11 +58,13 @@ class ExecutorStateReconciler:
         config_path: Optional[str],
         downloader: Downloader,
         task_reporter: TaskReporter,
-        server_channel: grpc.aio.Channel,
+        channel_creator: ChannelCreator,
         logger: Any,
     ):
         self._executor_id: str = executor_id
-        self._factory: FunctionExecutorServerFactory = function_executor_server_factory
+        self._function_executor_server_factory: FunctionExecutorServerFactory = (
+            function_executor_server_factory
+        )
         self._base_url: str = base_url
         self._config_path: Optional[str] = config_path
         self._downloader: Downloader = downloader
@@ -67,39 +72,60 @@ class ExecutorStateReconciler:
         self._function_executor_states: FunctionExecutorStatesContainer = (
             function_executor_states
         )
-        self._stub: TaskSchedulerServiceStub = TaskSchedulerServiceStub(server_channel)
+        self._channel_creator = channel_creator
         self._logger: Any = logger.bind(module=__name__)
         self._is_shutdown: bool = False
-        self._reconciliation_lock: asyncio.Lock = asyncio.Lock()
         self._server_last_clock: Optional[int] = None
 
     async def run(self):
-        desired_states: AsyncGenerator[DesiredExecutorState, None] = (
-            self._stub.get_desired_executor_states(
-                GetDesiredExecutorStatesRequest(executor_id=self._executor_id)
-            )
-        )
+        """Runs the state reconciler.
+
+        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
+
+        self._logger.info("State reconciler shutdown.")
+
+    async def _process_desired_states_stream(
+        self, desired_states: AsyncGenerator[DesiredExecutorState, None]
+    ):
         async for new_state in desired_states:
             if self._is_shutdown:
                 return
+
             new_state: DesiredExecutorState
             if self._server_last_clock is not None:
                 if self._server_last_clock >= new_state.clock:
                     continue  # Duplicate or outdated message state sent by Server.
 
             self._server_last_clock = new_state.clock
-            asyncio.create_task(self._reconcile_state(new_state))
+            await self._reconcile_state(new_state)
 
     async def _reconcile_state(self, new_state: DesiredExecutorState):
-        if self._is_shutdown:
-            return
-
-        # Simple non concurrent implementation for now for the PoC.
-        # Obtain this lock to force only a single coroutine doing the reconciliation.
-        async with self._reconciliation_lock:
-            await self._reconcile_function_executors(new_state)
-            # TODO
-            # await self._reconcile_task_allocations(new_state)
+        await self._reconcile_function_executors(new_state)
+        # TODO
+        # await self._reconcile_task_allocations(new_state)
 
     async def shutdown(self):
         """Shuts down the state reconciler.
@@ -121,6 +147,7 @@ class ExecutorStateReconciler:
                     graph_name=desired_function_executor.graph_name,
                     graph_version=desired_function_executor.graph_version,
                     function_name=desired_function_executor.function_name,
+                    image_uri=desired_function_executor.image_uri,
                 )
             )
 
@@ -203,13 +230,15 @@ class ExecutorStateReconciler:
             logger=logger,
         )
         function_executor: FunctionExecutor = FunctionExecutor(
-            server_factory=self._factory, logger=logger
+            server_factory=self._function_executor_server_factory, logger=logger
         )
         config: FunctionExecutorServerConfiguration = (
             FunctionExecutorServerConfiguration(
                 executor_id=self._executor_id,
                 function_executor_id=description.id,
+                namespace=description.namespace,
                 image_uri=description.image_uri,
+                secret_names=list(description.secret_names),
             )
         )
         initialize_request: InitializeRequest = InitializeRequest(

indexify/executor/grpc/state_reporter.py

@@ -0,0 +1,199 @@
+import asyncio
+from typing import Any, Dict, List, Optional
+
+import grpc
+
+from indexify.proto.task_scheduler_pb2 import (
+    AllowedFunction,
+    ExecutorState,
+    ExecutorStatus,
+    FunctionExecutorDescription,
+)
+from indexify.proto.task_scheduler_pb2 import (
+    FunctionExecutorState as FunctionExecutorStateProto,
+)
+from indexify.proto.task_scheduler_pb2 import (
+    FunctionExecutorStatus as FunctionExecutorStatusProto,
+)
+from indexify.proto.task_scheduler_pb2 import (
+    GPUModel,
+    GPUResources,
+    HostResources,
+    ReportExecutorStateRequest,
+)
+from indexify.proto.task_scheduler_pb2_grpc import (
+    TaskSchedulerServiceStub,
+)
+
+from ..api_objects import FunctionURI
+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 .metrics.state_reporter import (
+    metric_state_report_errors,
+    metric_state_report_latency,
+    metric_state_report_rpcs,
+)
+
+_REPORTING_INTERVAL_SEC = 5
+_REPORT_RPC_TIMEOUT_SEC = 5
+_REPORT_BACKOFF_ON_ERROR_SEC = 5
+
+
+class ExecutorStateReporter:
+    def __init__(
+        self,
+        executor_id: str,
+        development_mode: bool,
+        function_allowlist: Optional[List[FunctionURI]],
+        function_executor_states: FunctionExecutorStatesContainer,
+        channel_creator: ChannelCreator,
+        logger: Any,
+    ):
+        self._executor_id: str = executor_id
+        self._development_mode: bool = development_mode
+        self._function_executor_states: FunctionExecutorStatesContainer = (
+            function_executor_states
+        )
+        self._channel_creator = channel_creator
+        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
+        )
+
+    def update_executor_status(self, value: ExecutorStatus):
+        self._executor_status = value
+
+    async def run(self):
+        """Runs the state reporter.
+
+        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
+
+        self._logger.info("State reporter shutdown")
+
+    async def _report_state(self, stub: TaskSchedulerServiceStub):
+        with (
+            metric_state_report_errors.count_exceptions(),
+            metric_state_report_latency.time(),
+        ):
+            metric_state_report_rpcs.inc()
+            state = ExecutorState(
+                executor_id=self._executor_id,
+                development_mode=self._development_mode,
+                executor_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(),
+            )
+
+            await stub.report_executor_state(
+                ReportExecutorStateRequest(executor_state=state),
+                timeout=_REPORT_RPC_TIMEOUT_SEC,
+            )
+
+    async def _fetch_free_host_resources(self) -> HostResources:
+        # TODO: Implement host resource metrics reporting.
+        return HostResources(
+            cpu_count=0,
+            memory_bytes=0,
+            disk_bytes=0,
+            gpu=GPUResources(
+                count=0,
+                model=GPUModel.GPU_MODEL_UNKNOWN,
+            ),
+        )
+
+    async def _fetch_function_executor_states(self) -> List[FunctionExecutorStateProto]:
+        states = []
+
+        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
+                    ),
+                )
+            )
+
+        return states
+
+    async def shutdown(self):
+        """Shuts down the state reporter.
+
+        Never raises any exceptions.
+        """
+        self._is_shutdown = True
+
+
+def _to_grpc_allowed_functions(function_allowlist: Optional[List[FunctionURI]]):
+    if function_allowlist is None:
+        return []
+
+    allowed_functions: List[AllowedFunction] = []
+    for function_uri in function_allowlist:
+        function_uri: FunctionURI
+        allowed_function = AllowedFunction(
+            namespace=function_uri.namespace,
+            graph_name=function_uri.compute_graph,
+            function_name=function_uri.compute_fn,
+        )
+        if function_uri.version is not None:
+            allowed_function.graph_version = function_uri.version
+        allowed_functions.append(allowed_function)
+
+    return allowed_functions
+
+
+_STATUS_MAPPING: Dict[FunctionExecutorStatus, Any] = {
+    FunctionExecutorStatus.STARTING_UP: FunctionExecutorStatusProto.FUNCTION_EXECUTOR_STATUS_STARTING_UP,
+    FunctionExecutorStatus.STARTUP_FAILED_CUSTOMER_ERROR: FunctionExecutorStatusProto.FUNCTION_EXECUTOR_STATUS_STARTUP_FAILED_CUSTOMER_ERROR,
+    FunctionExecutorStatus.STARTUP_FAILED_PLATFORM_ERROR: FunctionExecutorStatusProto.FUNCTION_EXECUTOR_STATUS_STARTUP_FAILED_PLATFORM_ERROR,
+    FunctionExecutorStatus.IDLE: FunctionExecutorStatusProto.FUNCTION_EXECUTOR_STATUS_IDLE,
+    FunctionExecutorStatus.RUNNING_TASK: FunctionExecutorStatusProto.FUNCTION_EXECUTOR_STATUS_RUNNING_TASK,
+    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,
+}
+
+
+def _to_grpc_function_executor_status(
+    status: FunctionExecutorStatus, logger: Any
+) -> FunctionExecutorStatusProto:
+    result: FunctionExecutorStatusProto = _STATUS_MAPPING.get(
+        status, FunctionExecutorStatusProto.FUNCTION_EXECUTOR_STATUS_UNKNOWN
+    )
+
+    if result == FunctionExecutorStatusProto.FUNCTION_EXECUTOR_STATUS_UNKNOWN:
+        logger.error("Unexpected Function Executor status", status=status)
+
+    return result

indexify/executor/metrics/task_runner.py

@@ -23,6 +23,13 @@ metric_tasks_blocked_by_policy: prometheus_client.Gauge = prometheus_client.Gaug
     "tasks_blocked_by_policy",
     "Number of tasks that are ready for execution but are blocked according to the current policy (typically waiting for a free Function Executor)",
 )
+metric_tasks_blocked_by_policy_per_function_name: prometheus_client.Gauge = (
+    prometheus_client.Gauge(
+        "tasks_blocked_by_policy_per_function_name",
+        "Number of tasks that are ready for execution but are blocked according to the current policy (typically waiting for a free Function Executor)",
+        ["function_name"],
+    )
+)
 
 # Metrics for the stage when task is running.
 metric_task_runs: prometheus_client.Counter = prometheus_client.Counter(

indexify/executor/monitoring/health_checker/generic_health_checker.py

@@ -3,6 +3,7 @@ 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,6 +17,7 @@ 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
@@ -42,17 +44,30 @@ class GenericHealthChecker(HealthChecker):
         # * 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.health_check_failed:
-                return HealthCheckResult(
-                    is_success=False,
-                    status_message="A Function Executor health check failed",
-                    checker_name=HEALTH_CHECKER_NAME,
-                )
-
-        return HealthCheckResult(
-            is_success=True,
-            status_message="All Function Executors pass health checks",
-            checker_name=HEALTH_CHECKER_NAME,
-        )
+            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

indexify/executor/task_runner.py

@@ -1,7 +1,10 @@
 from typing import Any, Optional
 
 from .api_objects import Task
-from .function_executor.function_executor_state import FunctionExecutorState
+from .function_executor.function_executor_state import (
+    FunctionExecutorState,
+    FunctionExecutorStatus,
+)
 from .function_executor.function_executor_states_container import (
     FunctionExecutorStatesContainer,
 )
@@ -19,6 +22,7 @@ from .metrics.task_runner import (
     metric_task_run_platform_errors,
     metric_task_runs,
     metric_tasks_blocked_by_policy,
+    metric_tasks_blocked_by_policy_per_function_name,
     metric_tasks_running,
 )
 
@@ -52,6 +56,9 @@ class TaskRunner:
             with (
                 metric_task_policy_errors.count_exceptions(),
                 metric_tasks_blocked_by_policy.track_inprogress(),
+                metric_tasks_blocked_by_policy_per_function_name.labels(
+                    function_name=task_input.task.compute_fn
+                ).track_inprogress(),
                 metric_task_policy_latency.time(),
             ):
                 metric_task_policy_runs.inc()
@@ -109,6 +116,10 @@ class TaskRunner:
             raise
 
     async def _run_task_policy(self, state: FunctionExecutorState, task: Task) -> None:
+        """Runs the task policy until the task can run on the Function Executor.
+
+        On successful return the Function Executor status is either IDLE or DESTROYED.
+        """
         # Current policy for running tasks:
         #   - There can only be a single Function Executor per function regardless of function versions.
         #   --  If a Function Executor already exists for a different function version then wait until
@@ -116,13 +127,30 @@ class TaskRunner:
         #   --  This prevents failed tasks for different versions of the same function continiously
         #       destroying each other's Function Executors.
         #   - Each Function Executor rans at most 1 task concurrently.
-        await state.wait_running_tasks_less(1)
+        await state.wait_status(
+            [
+                FunctionExecutorStatus.DESTROYED,
+                FunctionExecutorStatus.IDLE,
+                FunctionExecutorStatus.UNHEALTHY,
+                FunctionExecutorStatus.SHUTDOWN,
+            ]
+        )
+        # We only shutdown the Function Executor on full Executor shutdown so it's fine to raise error here.
+        if state.status == FunctionExecutorStatus.SHUTDOWN:
+            raise Exception("Function Executor state is shutting down")
 
-        if state.graph_version != task.graph_version:
+        if state.status == FunctionExecutorStatus.UNHEALTHY:
             await state.destroy_function_executor()
-            state.graph_version = task.graph_version
-            # At this point the state belongs to the version of the function from the task
-            # and there are no running tasks in the Function Executor.
+
+        if state.graph_version == task.graph_version:
+            return  # All good, we can run on this Function Executor.
+
+        if state.status in [FunctionExecutorStatus.IDLE]:
+            await state.destroy_function_executor()
+
+        state.graph_version = task.graph_version
+        # At this point the state belongs to the version of the function from the task
+        # and there are no running tasks in the Function Executor.
 
     async def _run_task(
         self, state: FunctionExecutorState, task_input: TaskInput, logger: Any

indexify/{task_scheduler/proto → proto}/task_scheduler.proto

@@ -26,7 +26,7 @@ message GPUResources {
     optional GPUModel model = 2;
 }
 
-// Free host resources available at the Executor.
+// Resources that we're currently tracking and limiting on Executor.
 message HostResources {
     optional uint32 cpu_count = 1;
     optional uint64 memory_bytes = 2;
@@ -45,14 +45,14 @@ message AllowedFunction {
 
 enum FunctionExecutorStatus {
     FUNCTION_EXECUTOR_STATUS_UNKNOWN = 0;
-    FUNCTION_EXECUTOR_STATUS_STOPPED = 1;
-    FUNCTION_EXECUTOR_STATUS_STARTING_UP = 2;
-    FUNCTION_EXECUTOR_STATUS_STARTUP_FAILED_CUSTOMER_ERROR = 3;
-    FUNCTION_EXECUTOR_STATUS_STARTUP_FAILED_PLATFORM_ERROR = 4;
-    FUNCTION_EXECUTOR_STATUS_IDLE = 5;
-    FUNCTION_EXECUTOR_STATUS_RUNNING_TASK = 6;
-    FUNCTION_EXECUTOR_STATUS_UNHEALTHY = 7;
-    FUNCTION_EXECUTOR_STATUS_STOPPING = 8;
+    FUNCTION_EXECUTOR_STATUS_STARTING_UP = 1;
+    FUNCTION_EXECUTOR_STATUS_STARTUP_FAILED_CUSTOMER_ERROR = 2;
+    FUNCTION_EXECUTOR_STATUS_STARTUP_FAILED_PLATFORM_ERROR = 3;
+    FUNCTION_EXECUTOR_STATUS_IDLE = 4;
+    FUNCTION_EXECUTOR_STATUS_RUNNING_TASK = 5;
+    FUNCTION_EXECUTOR_STATUS_UNHEALTHY = 6;
+    FUNCTION_EXECUTOR_STATUS_STOPPING = 7;
+    FUNCTION_EXECUTOR_STATUS_STOPPED = 8;
 }
 
 // Immutable information that identifies and describes a Function Executor.
@@ -63,6 +63,8 @@ message FunctionExecutorDescription {
     optional string graph_version = 4;
     optional string function_name = 5;
     optional string image_uri = 6;
+    repeated string secret_names = 7;
+    optional HostResources resource_limits = 8;
 }
 
 message FunctionExecutorState {
@@ -72,19 +74,22 @@ message FunctionExecutorState {
 
 enum ExecutorStatus {
     EXECUTOR_STATUS_UNKNOWN = 0;
-    EXECUTOR_STATUS_STARTING = 1;
+    EXECUTOR_STATUS_STARTING_UP = 1;
     EXECUTOR_STATUS_RUNNING = 2;
     EXECUTOR_STATUS_DRAINED = 3;
-    EXECUTOR_STATUS_SHUTTING_DOWN = 4;
+    EXECUTOR_STATUS_STOPPING = 4;
+    EXECUTOR_STATUS_STOPPED = 5;
 }
 
 message ExecutorState {
     optional string executor_id = 1;
-    optional ExecutorStatus executor_status = 2;
-    optional HostResources host_resources = 3;
+    optional bool development_mode = 2;
+    optional ExecutorStatus executor_status = 3;
+    // Free resources available at the Executor.
+    optional HostResources free_resources = 4;
     // Empty allowed_functions list means that any function can run on the Executor.
-    repeated AllowedFunction allowed_functions = 4;
-    repeated FunctionExecutorState function_executor_states = 5;
+    repeated AllowedFunction allowed_functions = 5;
+    repeated FunctionExecutorState function_executor_states = 6;
 }
 
 // A message sent by Executor to report its up to date state to Server.
@@ -106,6 +111,7 @@ message Task {
     optional string graph_invocation_id = 6;
     optional string input_key = 8;
     optional string reducer_output_key = 9;
+    optional string timeout_ms = 10;
 }
 
 message TaskAllocation {
@@ -139,9 +145,9 @@ service TaskSchedulerService {
     // Called by Executor to open a stream of its desired states. When Server wants Executor to change something
     // it puts a message on the stream with the new desired state of the Executor.
     //
-    // Depricated HTTP API is used to download the serialized graph and task inputs.
+    // Deprecated HTTP API is used to download the serialized graph and task inputs.
     rpc get_desired_executor_states(GetDesiredExecutorStatesRequest) returns (stream DesiredExecutorState) {}
 
-    // Task outcome is currently reported via depricated HTTP API. We're going to migrate task output reporting to gRPC
+    // Task outcome is currently reported via deprecated HTTP API. We're going to migrate task output reporting to gRPC
     // when we move S3 downloads and uploads to Executor.
 }