indexify 0.3.14__py3-none-any.whl → 0.3.16__py3-none-any.whl

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/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,
 )
@@ -109,6 +112,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 +123,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.
 }

indexify/proto/task_scheduler_pb2.py

@@ -0,0 +1,64 @@
+# -*- coding: utf-8 -*-
+# Generated by the protocol buffer compiler.  DO NOT EDIT!
+# NO CHECKED-IN PROTOBUF GENCODE
+# source: indexify/proto/task_scheduler.proto
+# Protobuf Python Version: 5.29.0
+"""Generated protocol buffer code."""
+from google.protobuf import descriptor as _descriptor
+from google.protobuf import descriptor_pool as _descriptor_pool
+from google.protobuf import runtime_version as _runtime_version
+from google.protobuf import symbol_database as _symbol_database
+from google.protobuf.internal import builder as _builder
+
+_runtime_version.ValidateProtobufRuntimeVersion(
+    _runtime_version.Domain.PUBLIC, 5, 29, 0, "", "indexify/proto/task_scheduler.proto"
+)
+# @@protoc_insertion_point(imports)
+
+_sym_db = _symbol_database.Default()
+
+
+DESCRIPTOR = _descriptor_pool.Default().AddSerializedFile(
+    b'\n#indexify/proto/task_scheduler.proto\x12\x16task_scheduler_service"l\n\x0cGPUResources\x12\x12\n\x05\x63ount\x18\x01 \x01(\rH\x00\x88\x01\x01\x12\x34\n\x05model\x18\x02 \x01(\x0e\x32 .task_scheduler_service.GPUModelH\x01\x88\x01\x01\x42\x08\n\x06_countB\x08\n\x06_model"\xc9\x01\n\rHostResources\x12\x16\n\tcpu_count\x18\x01 \x01(\rH\x00\x88\x01\x01\x12\x19\n\x0cmemory_bytes\x18\x02 \x01(\x04H\x01\x88\x01\x01\x12\x17\n\ndisk_bytes\x18\x03 \x01(\x04H\x02\x88\x01\x01\x12\x36\n\x03gpu\x18\x04 \x01(\x0b\x32$.task_scheduler_service.GPUResourcesH\x03\x88\x01\x01\x42\x0c\n\n_cpu_countB\x0f\n\r_memory_bytesB\r\n\x0b_disk_bytesB\x06\n\x04_gpu"\xbb\x01\n\x0f\x41llowedFunction\x12\x16\n\tnamespace\x18\x01 \x01(\tH\x00\x88\x01\x01\x12\x17\n\ngraph_name\x18\x02 \x01(\tH\x01\x88\x01\x01\x12\x1a\n\rfunction_name\x18\x03 \x01(\tH\x02\x88\x01\x01\x12\x1a\n\rgraph_version\x18\x04 \x01(\tH\x03\x88\x01\x01\x42\x0c\n\n_namespaceB\r\n\x0b_graph_nameB\x10\n\x0e_function_nameB\x10\n\x0e_graph_version"\xf4\x02\n\x1b\x46unctionExecutorDescription\x12\x0f\n\x02id\x18\x01 \x01(\tH\x00\x88\x01\x01\x12\x16\n\tnamespace\x18\x02 \x01(\tH\x01\x88\x01\x01\x12\x17\n\ngraph_name\x18\x03 \x01(\tH\x02\x88\x01\x01\x12\x1a\n\rgraph_version\x18\x04 \x01(\tH\x03\x88\x01\x01\x12\x1a\n\rfunction_name\x18\x05 \x01(\tH\x04\x88\x01\x01\x12\x16\n\timage_uri\x18\x06 \x01(\tH\x05\x88\x01\x01\x12\x14\n\x0csecret_names\x18\x07 \x03(\t\x12\x43\n\x0fresource_limits\x18\x08 \x01(\x0b\x32%.task_scheduler_service.HostResourcesH\x06\x88\x01\x01\x42\x05\n\x03_idB\x0c\n\n_namespaceB\r\n\x0b_graph_nameB\x10\n\x0e_graph_versionB\x10\n\x0e_function_nameB\x0c\n\n_image_uriB\x12\n\x10_resource_limits"\xc6\x01\n\x15\x46unctionExecutorState\x12M\n\x0b\x64\x65scription\x18\x01 \x01(\x0b\x32\x33.task_scheduler_service.FunctionExecutorDescriptionH\x00\x88\x01\x01\x12\x43\n\x06status\x18\x02 \x01(\x0e\x32..task_scheduler_service.FunctionExecutorStatusH\x01\x88\x01\x01\x42\x0e\n\x0c_descriptionB\t\n\x07_status"\xb3\x03\n\rExecutorState\x12\x18\n\x0b\x65xecutor_id\x18\x01 \x01(\tH\x00\x88\x01\x01\x12\x1d\n\x10\x64\x65velopment_mode\x18\x02 \x01(\x08H\x01\x88\x01\x01\x12\x44\n\x0f\x65xecutor_status\x18\x03 \x01(\x0e\x32&.task_scheduler_service.ExecutorStatusH\x02\x88\x01\x01\x12\x42\n\x0e\x66ree_resources\x18\x04 \x01(\x0b\x32%.task_scheduler_service.HostResourcesH\x03\x88\x01\x01\x12\x42\n\x11\x61llowed_functions\x18\x05 \x03(\x0b\x32\'.task_scheduler_service.AllowedFunction\x12O\n\x18\x66unction_executor_states\x18\x06 \x03(\x0b\x32-.task_scheduler_service.FunctionExecutorStateB\x0e\n\x0c_executor_idB\x13\n\x11_development_modeB\x12\n\x10_executor_statusB\x11\n\x0f_free_resources"s\n\x1aReportExecutorStateRequest\x12\x42\n\x0e\x65xecutor_state\x18\x01 \x01(\x0b\x32%.task_scheduler_service.ExecutorStateH\x00\x88\x01\x01\x42\x11\n\x0f_executor_state"\x1d\n\x1bReportExecutorStateResponse"\x88\x03\n\x04Task\x12\x0f\n\x02id\x18\x01 \x01(\tH\x00\x88\x01\x01\x12\x16\n\tnamespace\x18\x02 \x01(\tH\x01\x88\x01\x01\x12\x17\n\ngraph_name\x18\x03 \x01(\tH\x02\x88\x01\x01\x12\x1a\n\rgraph_version\x18\x04 \x01(\tH\x03\x88\x01\x01\x12\x1a\n\rfunction_name\x18\x05 \x01(\tH\x04\x88\x01\x01\x12 \n\x13graph_invocation_id\x18\x06 \x01(\tH\x05\x88\x01\x01\x12\x16\n\tinput_key\x18\x08 \x01(\tH\x06\x88\x01\x01\x12\x1f\n\x12reducer_output_key\x18\t \x01(\tH\x07\x88\x01\x01\x12\x17\n\ntimeout_ms\x18\n \x01(\tH\x08\x88\x01\x01\x42\x05\n\x03_idB\x0c\n\n_namespaceB\r\n\x0b_graph_nameB\x10\n\x0e_graph_versionB\x10\n\x0e_function_nameB\x16\n\x14_graph_invocation_idB\x0c\n\n_input_keyB\x15\n\x13_reducer_output_keyB\r\n\x0b_timeout_ms"\x86\x01\n\x0eTaskAllocation\x12!\n\x14\x66unction_executor_id\x18\x01 \x01(\tH\x00\x88\x01\x01\x12/\n\x04task\x18\x02 \x01(\x0b\x32\x1c.task_scheduler_service.TaskH\x01\x88\x01\x01\x42\x17\n\x15_function_executor_idB\x07\n\x05_task"K\n\x1fGetDesiredExecutorStatesRequest\x12\x18\n\x0b\x65xecutor_id\x18\x01 \x01(\tH\x00\x88\x01\x01\x42\x0e\n\x0c_executor_id"\xc7\x01\n\x14\x44\x65siredExecutorState\x12O\n\x12\x66unction_executors\x18\x01 \x03(\x0b\x32\x33.task_scheduler_service.FunctionExecutorDescription\x12@\n\x10task_allocations\x18\x02 \x03(\x0b\x32&.task_scheduler_service.TaskAllocation\x12\x12\n\x05\x63lock\x18\x03 \x01(\x04H\x00\x88\x01\x01\x42\x08\n\x06_clock*\x86\x03\n\x08GPUModel\x12\x15\n\x11GPU_MODEL_UNKNOWN\x10\x00\x12"\n\x1eGPU_MODEL_NVIDIA_TESLA_T4_16GB\x10\n\x12$\n GPU_MODEL_NVIDIA_TESLA_V100_16GB\x10\x14\x12\x1d\n\x19GPU_MODEL_NVIDIA_A10_24GB\x10\x1e\x12\x1f\n\x1bGPU_MODEL_NVIDIA_A6000_48GB\x10(\x12#\n\x1fGPU_MODEL_NVIDIA_A100_SXM4_40GB\x10\x32\x12#\n\x1fGPU_MODEL_NVIDIA_A100_SXM4_80GB\x10\x33\x12"\n\x1eGPU_MODEL_NVIDIA_A100_PCI_40GB\x10\x34\x12#\n\x1fGPU_MODEL_NVIDIA_H100_SXM5_80GB\x10<\x12"\n\x1eGPU_MODEL_NVIDIA_H100_PCI_80GB\x10=\x12"\n\x1eGPU_MODEL_NVIDIA_RTX_6000_24GB\x10>*\xa3\x03\n\x16\x46unctionExecutorStatus\x12$\n FUNCTION_EXECUTOR_STATUS_UNKNOWN\x10\x00\x12(\n$FUNCTION_EXECUTOR_STATUS_STARTING_UP\x10\x01\x12:\n6FUNCTION_EXECUTOR_STATUS_STARTUP_FAILED_CUSTOMER_ERROR\x10\x02\x12:\n6FUNCTION_EXECUTOR_STATUS_STARTUP_FAILED_PLATFORM_ERROR\x10\x03\x12!\n\x1d\x46UNCTION_EXECUTOR_STATUS_IDLE\x10\x04\x12)\n%FUNCTION_EXECUTOR_STATUS_RUNNING_TASK\x10\x05\x12&\n"FUNCTION_EXECUTOR_STATUS_UNHEALTHY\x10\x06\x12%\n!FUNCTION_EXECUTOR_STATUS_STOPPING\x10\x07\x12$\n FUNCTION_EXECUTOR_STATUS_STOPPED\x10\x08*\xc3\x01\n\x0e\x45xecutorStatus\x12\x1b\n\x17\x45XECUTOR_STATUS_UNKNOWN\x10\x00\x12\x1f\n\x1b\x45XECUTOR_STATUS_STARTING_UP\x10\x01\x12\x1b\n\x17\x45XECUTOR_STATUS_RUNNING\x10\x02\x12\x1b\n\x17\x45XECUTOR_STATUS_DRAINED\x10\x03\x12\x1c\n\x18\x45XECUTOR_STATUS_STOPPING\x10\x04\x12\x1b\n\x17\x45XECUTOR_STATUS_STOPPED\x10\x05\x32\xa6\x02\n\x14TaskSchedulerService\x12\x82\x01\n\x15report_executor_state\x12\x32.task_scheduler_service.ReportExecutorStateRequest\x1a\x33.task_scheduler_service.ReportExecutorStateResponse"\x00\x12\x88\x01\n\x1bget_desired_executor_states\x12\x37.task_scheduler_service.GetDesiredExecutorStatesRequest\x1a,.task_scheduler_service.DesiredExecutorState"\x00\x30\x01\x62\x06proto3'
+)
+
+_globals = globals()
+_builder.BuildMessageAndEnumDescriptors(DESCRIPTOR, _globals)
+_builder.BuildTopDescriptorsAndMessages(
+    DESCRIPTOR, "indexify.proto.task_scheduler_pb2", _globals
+)
+if not _descriptor._USE_C_DESCRIPTORS:
+    DESCRIPTOR._loaded_options = None
+    _globals["_GPUMODEL"]._serialized_start = 2541
+    _globals["_GPUMODEL"]._serialized_end = 2931
+    _globals["_FUNCTIONEXECUTORSTATUS"]._serialized_start = 2934
+    _globals["_FUNCTIONEXECUTORSTATUS"]._serialized_end = 3353
+    _globals["_EXECUTORSTATUS"]._serialized_start = 3356
+    _globals["_EXECUTORSTATUS"]._serialized_end = 3551
+    _globals["_GPURESOURCES"]._serialized_start = 63
+    _globals["_GPURESOURCES"]._serialized_end = 171
+    _globals["_HOSTRESOURCES"]._serialized_start = 174
+    _globals["_HOSTRESOURCES"]._serialized_end = 375
+    _globals["_ALLOWEDFUNCTION"]._serialized_start = 378
+    _globals["_ALLOWEDFUNCTION"]._serialized_end = 565
+    _globals["_FUNCTIONEXECUTORDESCRIPTION"]._serialized_start = 568
+    _globals["_FUNCTIONEXECUTORDESCRIPTION"]._serialized_end = 940
+    _globals["_FUNCTIONEXECUTORSTATE"]._serialized_start = 943
+    _globals["_FUNCTIONEXECUTORSTATE"]._serialized_end = 1141
+    _globals["_EXECUTORSTATE"]._serialized_start = 1144
+    _globals["_EXECUTORSTATE"]._serialized_end = 1579
+    _globals["_REPORTEXECUTORSTATEREQUEST"]._serialized_start = 1581
+    _globals["_REPORTEXECUTORSTATEREQUEST"]._serialized_end = 1696
+    _globals["_REPORTEXECUTORSTATERESPONSE"]._serialized_start = 1698
+    _globals["_REPORTEXECUTORSTATERESPONSE"]._serialized_end = 1727
+    _globals["_TASK"]._serialized_start = 1730
+    _globals["_TASK"]._serialized_end = 2122
+    _globals["_TASKALLOCATION"]._serialized_start = 2125
+    _globals["_TASKALLOCATION"]._serialized_end = 2259
+    _globals["_GETDESIREDEXECUTORSTATESREQUEST"]._serialized_start = 2261
+    _globals["_GETDESIREDEXECUTORSTATESREQUEST"]._serialized_end = 2336
+    _globals["_DESIREDEXECUTORSTATE"]._serialized_start = 2339
+    _globals["_DESIREDEXECUTORSTATE"]._serialized_end = 2538
+    _globals["_TASKSCHEDULERSERVICE"]._serialized_start = 3554
+    _globals["_TASKSCHEDULERSERVICE"]._serialized_end = 3848
+# @@protoc_insertion_point(module_scope)

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

@@ -28,7 +28,6 @@ class GPUModel(int, metaclass=_enum_type_wrapper.EnumTypeWrapper):
 class FunctionExecutorStatus(int, metaclass=_enum_type_wrapper.EnumTypeWrapper):
     __slots__ = ()
     FUNCTION_EXECUTOR_STATUS_UNKNOWN: _ClassVar[FunctionExecutorStatus]
-    FUNCTION_EXECUTOR_STATUS_STOPPED: _ClassVar[FunctionExecutorStatus]
     FUNCTION_EXECUTOR_STATUS_STARTING_UP: _ClassVar[FunctionExecutorStatus]
     FUNCTION_EXECUTOR_STATUS_STARTUP_FAILED_CUSTOMER_ERROR: _ClassVar[
         FunctionExecutorStatus
@@ -40,14 +39,16 @@ class FunctionExecutorStatus(int, metaclass=_enum_type_wrapper.EnumTypeWrapper):
     FUNCTION_EXECUTOR_STATUS_RUNNING_TASK: _ClassVar[FunctionExecutorStatus]
     FUNCTION_EXECUTOR_STATUS_UNHEALTHY: _ClassVar[FunctionExecutorStatus]
     FUNCTION_EXECUTOR_STATUS_STOPPING: _ClassVar[FunctionExecutorStatus]
+    FUNCTION_EXECUTOR_STATUS_STOPPED: _ClassVar[FunctionExecutorStatus]
 
 class ExecutorStatus(int, metaclass=_enum_type_wrapper.EnumTypeWrapper):
     __slots__ = ()
     EXECUTOR_STATUS_UNKNOWN: _ClassVar[ExecutorStatus]
-    EXECUTOR_STATUS_STARTING: _ClassVar[ExecutorStatus]
+    EXECUTOR_STATUS_STARTING_UP: _ClassVar[ExecutorStatus]
     EXECUTOR_STATUS_RUNNING: _ClassVar[ExecutorStatus]
     EXECUTOR_STATUS_DRAINED: _ClassVar[ExecutorStatus]
-    EXECUTOR_STATUS_SHUTTING_DOWN: _ClassVar[ExecutorStatus]
+    EXECUTOR_STATUS_STOPPING: _ClassVar[ExecutorStatus]
+    EXECUTOR_STATUS_STOPPED: _ClassVar[ExecutorStatus]
 
 GPU_MODEL_UNKNOWN: GPUModel
 GPU_MODEL_NVIDIA_TESLA_T4_16GB: GPUModel
@@ -61,7 +62,6 @@ GPU_MODEL_NVIDIA_H100_SXM5_80GB: GPUModel
 GPU_MODEL_NVIDIA_H100_PCI_80GB: GPUModel
 GPU_MODEL_NVIDIA_RTX_6000_24GB: GPUModel
 FUNCTION_EXECUTOR_STATUS_UNKNOWN: FunctionExecutorStatus
-FUNCTION_EXECUTOR_STATUS_STOPPED: FunctionExecutorStatus
 FUNCTION_EXECUTOR_STATUS_STARTING_UP: FunctionExecutorStatus
 FUNCTION_EXECUTOR_STATUS_STARTUP_FAILED_CUSTOMER_ERROR: FunctionExecutorStatus
 FUNCTION_EXECUTOR_STATUS_STARTUP_FAILED_PLATFORM_ERROR: FunctionExecutorStatus
@@ -69,11 +69,13 @@ FUNCTION_EXECUTOR_STATUS_IDLE: FunctionExecutorStatus
 FUNCTION_EXECUTOR_STATUS_RUNNING_TASK: FunctionExecutorStatus
 FUNCTION_EXECUTOR_STATUS_UNHEALTHY: FunctionExecutorStatus
 FUNCTION_EXECUTOR_STATUS_STOPPING: FunctionExecutorStatus
+FUNCTION_EXECUTOR_STATUS_STOPPED: FunctionExecutorStatus
 EXECUTOR_STATUS_UNKNOWN: ExecutorStatus
-EXECUTOR_STATUS_STARTING: ExecutorStatus
+EXECUTOR_STATUS_STARTING_UP: ExecutorStatus
 EXECUTOR_STATUS_RUNNING: ExecutorStatus
 EXECUTOR_STATUS_DRAINED: ExecutorStatus
-EXECUTOR_STATUS_SHUTTING_DOWN: ExecutorStatus
+EXECUTOR_STATUS_STOPPING: ExecutorStatus
+EXECUTOR_STATUS_STOPPED: ExecutorStatus
 
 class GPUResources(_message.Message):
     __slots__ = ("count", "model")
@@ -129,6 +131,8 @@ class FunctionExecutorDescription(_message.Message):
         "graph_version",
         "function_name",
         "image_uri",
+        "secret_names",
+        "resource_limits",
     )
     ID_FIELD_NUMBER: _ClassVar[int]
     NAMESPACE_FIELD_NUMBER: _ClassVar[int]
@@ -136,12 +140,16 @@ class FunctionExecutorDescription(_message.Message):
     GRAPH_VERSION_FIELD_NUMBER: _ClassVar[int]
     FUNCTION_NAME_FIELD_NUMBER: _ClassVar[int]
     IMAGE_URI_FIELD_NUMBER: _ClassVar[int]
+    SECRET_NAMES_FIELD_NUMBER: _ClassVar[int]
+    RESOURCE_LIMITS_FIELD_NUMBER: _ClassVar[int]
     id: str
     namespace: str
     graph_name: str
     graph_version: str
     function_name: str
     image_uri: str
+    secret_names: _containers.RepeatedScalarFieldContainer[str]
+    resource_limits: HostResources
     def __init__(
         self,
         id: _Optional[str] = ...,
@@ -150,6 +158,8 @@ class FunctionExecutorDescription(_message.Message):
         graph_version: _Optional[str] = ...,
         function_name: _Optional[str] = ...,
         image_uri: _Optional[str] = ...,
+        secret_names: _Optional[_Iterable[str]] = ...,
+        resource_limits: _Optional[_Union[HostResources, _Mapping]] = ...,
     ) -> None: ...
 
 class FunctionExecutorState(_message.Message):
@@ -167,19 +177,22 @@ class FunctionExecutorState(_message.Message):
 class ExecutorState(_message.Message):
     __slots__ = (
         "executor_id",
+        "development_mode",
         "executor_status",
-        "host_resources",
+        "free_resources",
         "allowed_functions",
         "function_executor_states",
     )
     EXECUTOR_ID_FIELD_NUMBER: _ClassVar[int]
+    DEVELOPMENT_MODE_FIELD_NUMBER: _ClassVar[int]
     EXECUTOR_STATUS_FIELD_NUMBER: _ClassVar[int]
-    HOST_RESOURCES_FIELD_NUMBER: _ClassVar[int]
+    FREE_RESOURCES_FIELD_NUMBER: _ClassVar[int]
     ALLOWED_FUNCTIONS_FIELD_NUMBER: _ClassVar[int]
     FUNCTION_EXECUTOR_STATES_FIELD_NUMBER: _ClassVar[int]
     executor_id: str
+    development_mode: bool
     executor_status: ExecutorStatus
-    host_resources: HostResources
+    free_resources: HostResources
     allowed_functions: _containers.RepeatedCompositeFieldContainer[AllowedFunction]
     function_executor_states: _containers.RepeatedCompositeFieldContainer[
         FunctionExecutorState
@@ -187,8 +200,9 @@ class ExecutorState(_message.Message):
     def __init__(
         self,
         executor_id: _Optional[str] = ...,
+        development_mode: bool = ...,
         executor_status: _Optional[_Union[ExecutorStatus, str]] = ...,
-        host_resources: _Optional[_Union[HostResources, _Mapping]] = ...,
+        free_resources: _Optional[_Union[HostResources, _Mapping]] = ...,
         allowed_functions: _Optional[
             _Iterable[_Union[AllowedFunction, _Mapping]]
         ] = ...,
@@ -219,6 +233,7 @@ class Task(_message.Message):
         "graph_invocation_id",
         "input_key",
         "reducer_output_key",
+        "timeout_ms",
     )
     ID_FIELD_NUMBER: _ClassVar[int]
     NAMESPACE_FIELD_NUMBER: _ClassVar[int]
@@ -228,6 +243,7 @@ class Task(_message.Message):
     GRAPH_INVOCATION_ID_FIELD_NUMBER: _ClassVar[int]
     INPUT_KEY_FIELD_NUMBER: _ClassVar[int]
     REDUCER_OUTPUT_KEY_FIELD_NUMBER: _ClassVar[int]
+    TIMEOUT_MS_FIELD_NUMBER: _ClassVar[int]
     id: str
     namespace: str
     graph_name: str
@@ -236,6 +252,7 @@ class Task(_message.Message):
     graph_invocation_id: str
     input_key: str
     reducer_output_key: str
+    timeout_ms: str
     def __init__(
         self,
         id: _Optional[str] = ...,
@@ -246,6 +263,7 @@ class Task(_message.Message):
         graph_invocation_id: _Optional[str] = ...,
         input_key: _Optional[str] = ...,
         reducer_output_key: _Optional[str] = ...,
+        timeout_ms: _Optional[str] = ...,
     ) -> None: ...
 
 class TaskAllocation(_message.Message):