indexify 0.3.18__tar.gz → 0.3.19__tar.gz

{indexify-0.3.18 → indexify-0.3.19}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.3
 Name: indexify
-Version: 0.3.18
+Version: 0.3.19
 Summary: Open Source Indexify components and helper tools
 Home-page: https://github.com/tensorlakeai/indexify
 License: Apache 2.0

{indexify-0.3.18 → indexify-0.3.19}/pyproject.toml

@@ -1,7 +1,7 @@
 [tool.poetry]
 name = "indexify"
 # Incremented if any of the components provided in this packages are updated.
-version = "0.3.18"
+version = "0.3.19"
 description = "Open Source Indexify components and helper tools"
 authors = ["Tensorlake Inc. <support@tensorlake.ai>"]
 license = "Apache 2.0"
@@ -24,6 +24,8 @@ aiohttp = "^3.11.0"
 prometheus-client = "^0.21.1"
 # Adds function-executor binary and utils lib.
 tensorlake = ">=0.1"
+# Uncomment the next line to use local tensorlake package (only for development!)
+# tensorlake = { path = "../tensorlake", develop = true }
 # pydantic is provided by tensorlake
 # httpx-sse is provided by tensorlake
 # grpcio is provided by tensorlake

{indexify-0.3.18 → indexify-0.3.19}/src/indexify/cli/cli.py

@@ -78,6 +78,7 @@ def build_image(
 )
 def executor(
     server_addr: str = "localhost:8900",
+    grpc_server_addr: str = "localhost:8901",
     dev: Annotated[
         bool, typer.Option("--dev", "-d", help="Run the executor in development mode")
     ] = False,
@@ -120,16 +121,6 @@ def executor(
             help="Port where to run Executor Monitoring server",
         ),
     ] = 7000,
-    grpc_server_addr: Annotated[
-        Optional[str],
-        typer.Option(
-            "--grpc-server-addr",
-            help=(
-                "(exprimental) Address of server gRPC API to connect to, e.g. 'localhost:8901'.\n"
-                "Enables gRPC state reporter that will periodically report the state of the Function Executors to Server\n"
-            ),
-        ),
-    ] = None,
     enable_grpc_state_reconciler: Annotated[
         bool,
         typer.Option(
@@ -166,11 +157,6 @@ def executor(
             "--executor-id should be at least 10 characters long and only include characters _-[0-9][a-z][A-Z]"
         )
 
-    if enable_grpc_state_reconciler and grpc_server_addr is None:
-        raise typer.BadParameter(
-            "--grpc-server-addr must be set when --enable-grpc-state-reconciler is set"
-        )
-
     kv_labels: Dict[str, str] = {}
     for label in labels:
         key, value = label.split("=")
@@ -183,6 +169,7 @@ def executor(
         "starting executor",
         hostname=gethostname(),
         server_addr=server_addr,
+        grpc_server_addr=grpc_server_addr,
         config_path=config_path,
         executor_version=executor_version,
         labels=kv_labels,
@@ -192,7 +179,6 @@ def executor(
         dev_mode=dev,
         monitoring_server_host=monitoring_server_host,
         monitoring_server_port=monitoring_server_port,
-        grpc_server_addr=grpc_server_addr,
         enable_grpc_state_reconciler=enable_grpc_state_reconciler,
     )
 
@@ -231,10 +217,10 @@ def executor(
             server_ports=range(ports[0], ports[1]),
         ),
         server_addr=server_addr,
+        grpc_server_addr=grpc_server_addr,
         config_path=config_path,
         monitoring_server_host=monitoring_server_host,
         monitoring_server_port=monitoring_server_port,
-        grpc_server_addr=grpc_server_addr,
         enable_grpc_state_reconciler=enable_grpc_state_reconciler,
     ).run()
 

{indexify-0.3.18 → indexify-0.3.19}/src/indexify/executor/api_objects.py

@@ -49,5 +49,17 @@ class TaskResult(BaseModel):
     reducer: bool = False
 
 
+class DataPayload(BaseModel):
+    path: str
+    size: int
+    sha256_hash: str
+
+
+class IngestFnOutputsResponse(BaseModel):
+    data_payloads: List[DataPayload]
+    stdout: Optional[DataPayload] = None
+    stderr: Optional[DataPayload] = None
+
+
 TASK_OUTCOME_SUCCESS = "success"
 TASK_OUTCOME_FAILURE = "failure"

{indexify-0.3.18 → indexify-0.3.19}/src/indexify/executor/downloader.py

@@ -241,7 +241,10 @@ class Downloader:
 def serialized_object_from_http_response(response: httpx.Response) -> SerializedObject:
     # We're hardcoding the content type currently used by Python SDK. It might change in the future.
     # There's no other way for now to determine if the response is a bytes or string.
-    if response.headers["content-type"] == "application/octet-stream":
+    if response.headers["content-type"] in [
+        "application/octet-stream",
+        "application/pickle",
+    ]:
         return SerializedObject(
             bytes=response.content, content_type=response.headers["content-type"]
         )

{indexify-0.3.18 → indexify-0.3.19}/src/indexify/executor/executor.py

@@ -64,10 +64,10 @@ 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,
     ):
         self._logger = structlog.get_logger(module=__name__)
@@ -97,43 +97,40 @@ class Executor:
         self._downloader = Downloader(
             code_path=code_path, base_url=self._base_url, 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 +168,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(),
         }
@@ -256,6 +253,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,6 +265,28 @@ 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,

{indexify-0.3.18 → indexify-0.3.19}/src/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 {customer_code_timeout_sec} sec expired"
+                ) from e
+            raise

{indexify-0.3.18 → indexify-0.3.19}/src/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-0.3.18 → indexify-0.3.19}/src/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

{indexify-0.3.18 → indexify-0.3.19}/src/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"
@@ -69,6 +70,7 @@ def is_status_change_allowed(
         ],
         FunctionExecutorStatus.RUNNING_TASK: [
             FunctionExecutorStatus.RUNNING_TASK,
+            FunctionExecutorStatus.DESTROYING,
             FunctionExecutorStatus.IDLE,
             FunctionExecutorStatus.UNHEALTHY,
             FunctionExecutorStatus.SHUTDOWN,

{indexify-0.3.18 → indexify-0.3.19}/src/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-0.3.18 → indexify-0.3.19}/src/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-0.3.18 → indexify-0.3.19}/src/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,26 @@ 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,
+    ) -> "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="Function execution timed out.",
+            is_internal_error=False,
+        )

indexify-0.3.19/src/indexify/executor/grpc/completed_tasks_container.py

@@ -0,0 +1,26 @@
+import asyncio
+from typing import List, Set
+
+
+class CompletedTasksContainer:
+    """An asyncio concurrent container for the completed task IDs."""
+
+    def __init__(self):
+        # The fields below are protected by the lock.
+        self._lock: asyncio.Lock = asyncio.Lock()
+        self._completed_task_ids: Set[str] = set()
+
+    async def add(self, task_id: str) -> None:
+        """Add a task to the container."""
+        async with self._lock:
+            self._completed_task_ids.add(task_id)
+
+    async def contains(self, task_id: str) -> bool:
+        """Check if the task is in the container."""
+        async with self._lock:
+            return task_id in self._completed_task_ids
+
+    async def replace(self, task_ids: List[str]) -> None:
+        """Replaces the task IDs with the supplied task IDs."""
+        async with self._lock:
+            self._completed_task_ids = set(task_ids)