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

indexify/cli/cli.py

@@ -9,11 +9,7 @@ configure_logging_early()
 import os
 import re
 import shutil
-import signal
-import subprocess
 import sys
-import threading
-import time
 from importlib.metadata import version
 from pathlib import Path
 from socket import gethostname
@@ -51,91 +47,6 @@ console = Console(theme=custom_theme)
 app = typer.Typer(pretty_exceptions_enable=False, no_args_is_help=True)
 
 
-@app.command(
-    help="Run server and executor in dev mode (Not recommended for production.)"
-)
-def server_dev_mode():
-    indexify_server_path = os.path.expanduser("~/.indexify/indexify-server")
-    if not os.path.exists(indexify_server_path):
-        print("indexify-server not found. Downloading...")
-        try:
-            download_command = subprocess.check_output(
-                ["curl", "-s", "https://getindexify.ai"], universal_newlines=True
-            )
-            subprocess.run(download_command, shell=True, check=True)
-        except subprocess.CalledProcessError as e:
-            print(f"failed to download indexify-server: {e}")
-            exit(1)
-        try:
-            os.makedirs(os.path.dirname(indexify_server_path), exist_ok=True)
-            shutil.move("indexify-server", indexify_server_path)
-        except Exception as e:
-            print(f"failed to move indexify-server to {indexify_server_path}: {e}")
-            exit(1)
-    print("starting indexify server and executor in dev mode...")
-    print("press Ctrl+C to stop the server and executor.")
-    print(f"server binary path: {indexify_server_path}")
-    commands: List[List[str]] = [
-        [indexify_server_path, "--dev"],
-        ["indexify-cli", "executor", "--dev"],
-    ]
-    processes = []
-    stop_event = threading.Event()
-
-    def handle_output(process):
-        for line in iter(process.stdout.readline, ""):
-            sys.stdout.write(line)
-            sys.stdout.flush()
-
-    def terminate_processes():
-        print("Terminating processes...")
-        stop_event.set()
-        for process in processes:
-            if process.poll() is None:
-                try:
-                    process.terminate()
-                    process.wait(timeout=5)
-                except subprocess.TimeoutExpired:
-                    print(f"Force killing process {process.pid}")
-                    process.kill()
-
-    def signal_handler(sig, frame):
-        print("\nCtrl+C pressed. Shutting down...")
-        terminate_processes()
-        sys.exit(0)
-
-    signal.signal(signal.SIGINT, signal_handler)
-    signal.signal(signal.SIGTERM, signal_handler)
-
-    for cmd in commands:
-        process = subprocess.Popen(
-            cmd,
-            stdout=subprocess.PIPE,
-            stderr=subprocess.STDOUT,
-            bufsize=1,
-            universal_newlines=True,
-            preexec_fn=os.setsid if os.name != "nt" else None,
-        )
-        processes.append(process)
-
-        thread = threading.Thread(target=handle_output, args=(process,))
-        thread.daemon = True
-        thread.start()
-
-    try:
-        while True:
-            time.sleep(1)
-            if all(process.poll() is not None for process in processes):
-                print("All processes have finished.")
-                break
-    except KeyboardInterrupt:
-        signal_handler(None, None)
-    finally:
-        terminate_processes()
-
-    print("Script execution completed.")
-
-
 @app.command(help="Build image for function names")
 def build_image(
     workflow_file_path: str,
@@ -208,17 +119,27 @@ def executor(
             help="Port where to run Executor Monitoring server",
         ),
     ] = 7000,
+    # TODO: Figure out mTLS for gRPC.
     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"
-                "If set disables automatic Function Executor management on Executor and uses the Server gRPC API\n"
-                "for Function Executor management and placement of tasks on them."
+                "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(
+            "--enable-grpc-state-reconciler",
+            help=(
+                "(exprimental) Enable gRPC state reconciler that will reconcile the state of the Function Executors and Task Allocations\n"
+                "with the desired state provided by Server. Required --grpc-server-addr to be set."
+            ),
+        ),
+    ] = False,
 ):
     if dev:
         configure_development_mode_logging()
@@ -236,6 +157,11 @@ 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"
+        )
+
     executor_version = version("indexify")
     logger = structlog.get_logger(module=__name__, executor_id=executor_id)
 
@@ -252,6 +178,7 @@ def executor(
         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,
     )
 
     executor_cache = Path(executor_cache).expanduser().absolute()
@@ -277,6 +204,7 @@ def executor(
 
     Executor(
         id=executor_id,
+        development_mode=dev,
         version=executor_version,
         health_checker=GenericHealthChecker(),
         code_path=executor_cache,
@@ -290,6 +218,7 @@ def executor(
         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/executor/api_objects.py

@@ -14,6 +14,8 @@ class Task(BaseModel):
     graph_version: str
     image_uri: Optional[str] = None
     "image_uri defines the URI of the image of this task. Optional since some executors do not require it."
+    secret_names: Optional[List[str]] = None
+    "secret_names defines the names of the secrets to set on function executor. Optional for backward compatibility."
 
 
 class FunctionURI(BaseModel):

indexify/executor/executor.py

@@ -5,11 +5,12 @@ from pathlib import Path
 from socket import gethostname
 from typing import Any, Dict, List, Optional
 
-import grpc
 import structlog
 from tensorlake.function_executor.proto.function_executor_pb2 import SerializedObject
 from tensorlake.utils.logging import suppress as suppress_logging
 
+from indexify.proto.task_scheduler_pb2 import ExecutorStatus
+
 from .api_objects import FunctionURI, Task
 from .downloader import Downloader
 from .function_executor.function_executor_states_container import (
@@ -18,6 +19,9 @@ from .function_executor.function_executor_states_container import (
 from .function_executor.server.function_executor_server_factory import (
     FunctionExecutorServerFactory,
 )
+from .grpc.channel_creator import ChannelCreator
+from .grpc.state_reconciler import ExecutorStateReconciler
+from .grpc.state_reporter import ExecutorStateReporter
 from .metrics.executor import (
     METRIC_TASKS_COMPLETED_OUTCOME_ALL,
     METRIC_TASKS_COMPLETED_OUTCOME_ERROR_CUSTOMER_CODE,
@@ -39,14 +43,10 @@ from .monitoring.health_checker.health_checker import HealthChecker
 from .monitoring.prometheus_metrics_handler import PrometheusMetricsHandler
 from .monitoring.server import MonitoringServer
 from .monitoring.startup_probe_handler import StartupProbeHandler
-from .state_reconciler import ExecutorStateReconciler
-from .state_reporter import ExecutorStateReporter
 from .task_fetcher import TaskFetcher
 from .task_reporter import TaskReporter
 from .task_runner import TaskInput, TaskOutput, TaskRunner
 
-EXECUTOR_GRPC_SERVER_READY_TIMEOUT_SEC = 10
-
 metric_executor_state.state("starting")
 
 
@@ -54,6 +54,7 @@ class Executor:
     def __init__(
         self,
         id: str,
+        development_mode: bool,
         version: str,
         code_path: Path,
         health_checker: HealthChecker,
@@ -64,10 +65,10 @@ class Executor:
         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__)
         self._is_shutdown: bool = False
-        self._config_path = config_path
         protocol: str = "http"
         if config_path:
             self._logger.info("running the extractor with TLS enabled")
@@ -84,7 +85,9 @@ class Executor:
             health_probe_handler=HealthCheckHandler(health_checker),
             metrics_handler=PrometheusMetricsHandler(),
         )
-        self._function_executor_states = FunctionExecutorStatesContainer()
+        self._function_executor_states = FunctionExecutorStatesContainer(
+            logger=self._logger
+        )
         health_checker.set_function_executor_states_container(
             self._function_executor_states
         )
@@ -94,24 +97,54 @@ class Executor:
         self._task_reporter = TaskReporter(
             base_url=self._base_url,
             executor_id=id,
-            config_path=self._config_path,
+            config_path=config_path,
         )
-        self._grpc_server_addr: Optional[str] = grpc_server_addr
-        self._id = id
         self._function_allowlist: Optional[List[FunctionURI]] = function_allowlist
         self._function_executor_server_factory = function_executor_server_factory
+
+        # HTTP mode services
+        self._task_runner: Optional[TaskRunner] = None
+        self._task_fetcher: Optional[TaskFetcher] = None
+        # gRPC mode services
+        self._channel_creator: Optional[ChannelCreator] = None
         self._state_reporter: Optional[ExecutorStateReporter] = None
         self._state_reconciler: Optional[ExecutorStateReconciler] = None
 
-        if self._grpc_server_addr is None:
-            self._task_runner: Optional[TaskRunner] = TaskRunner(
+        if grpc_server_addr is not None:
+            self._channel_creator = ChannelCreator(grpc_server_addr, self._logger)
+            self._state_reporter = ExecutorStateReporter(
+                executor_id=id,
+                development_mode=development_mode,
+                function_allowlist=self._function_allowlist,
+                function_executor_states=self._function_executor_states,
+                channel_creator=self._channel_creator,
+                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,
+                function_executor_server_factory=self._function_executor_server_factory,
+                base_url=self._base_url,
+                function_executor_states=self._function_executor_states,
+                config_path=config_path,
+                downloader=self._downloader,
+                task_reporter=self._task_reporter,
+                channel_creator=self._channel_creator,
+                logger=self._logger,
+            )
+        else:
+            self._task_runner = TaskRunner(
                 executor_id=id,
                 function_executor_server_factory=function_executor_server_factory,
                 base_url=self._base_url,
                 function_executor_states=self._function_executor_states,
                 config_path=config_path,
             )
-            self._task_fetcher: Optional[TaskFetcher] = TaskFetcher(
+            self._task_fetcher = TaskFetcher(
                 executor_id=id,
                 executor_version=version,
                 function_allowlist=function_allowlist,
@@ -122,11 +155,13 @@ class Executor:
 
         executor_info: Dict[str, str] = {
             "id": id,
+            "dev_mode": str(development_mode),
             "version": version,
             "code_path": str(code_path),
             "server_addr": server_addr,
             "config_path": str(config_path),
             "grpc_server_addr": str(grpc_server_addr),
+            "enable_grpc_state_reconciler": str(enable_grpc_state_reconciler),
             "hostname": gethostname(),
         }
         executor_info.update(function_allowlist_to_info_dict(function_allowlist))
@@ -146,84 +181,35 @@ class Executor:
             )
 
         asyncio.get_event_loop().create_task(self._monitoring_server.run())
+        if self._state_reporter is not None:
+            self._state_reporter.update_executor_status(
+                ExecutorStatus.EXECUTOR_STATUS_RUNNING
+            )
+            asyncio.get_event_loop().create_task(self._state_reporter.run())
 
-        try:
-            if self._grpc_server_addr is None:
-                asyncio.get_event_loop().run_until_complete(self._http_mode_loop())
-            else:
-                asyncio.get_event_loop().run_until_complete(self._grpc_mode_loop())
-        except asyncio.CancelledError:
-            pass  # Suppress this expected exception and return without error (normally).
-
-    async def _grpc_mode_loop(self):
         metric_executor_state.state("running")
         self._startup_probe_handler.set_ready()
 
-        while not self._is_shutdown:
-            async with self._establish_grpc_server_channel() as server_channel:
-                server_channel: grpc.aio.Channel
-                await self._run_grpc_mode_services(server_channel)
-                self._logger.warning(
-                    "grpc mode services exited, retrying in 5 seconds",
-                )
-                await asyncio.sleep(5)
-
-    async def _establish_grpc_server_channel(self) -> grpc.aio.Channel:
         try:
-            channel = grpc.aio.insecure_channel(self._grpc_server_addr)
-            await asyncio.wait_for(
-                channel.channel_ready(),
-                timeout=EXECUTOR_GRPC_SERVER_READY_TIMEOUT_SEC,
-            )
-            return channel
-        except Exception as e:
-            self._logger.error("failed establishing grpc server channel", exc_info=e)
-            raise
+            if self._state_reconciler is None:
+                asyncio.get_event_loop().run_until_complete(
+                    self._http_task_runner_loop()
+                )
+            else:
+                asyncio.get_event_loop().run_until_complete(
+                    self._grpc_state_reconciler_loop()
+                )
+        except asyncio.CancelledError:
+            pass  # Suppress this expected exception and return without error (normally).
 
-    async def _run_grpc_mode_services(self, server_channel: grpc.aio.Channel):
-        """Runs the gRPC mode services.
+    async def _grpc_state_reconciler_loop(self):
+        """Runs the gRPC state reconciler and state reporter.
 
         Never raises any exceptions."""
-        try:
-            self._state_reporter = ExecutorStateReporter(
-                executor_id=self._id,
-                function_allowlist=self._function_allowlist,
-                function_executor_states=self._function_executor_states,
-                server_channel=server_channel,
-                logger=self._logger,
-            )
-            self._state_reconciler = ExecutorStateReconciler(
-                executor_id=self._id,
-                function_executor_server_factory=self._function_executor_server_factory,
-                base_url=self._base_url,
-                function_executor_states=self._function_executor_states,
-                config_path=self._config_path,
-                downloader=self._downloader,
-                task_reporter=self._task_reporter,
-                server_channel=server_channel,
-                logger=self._logger,
-            )
+        asyncio.create_task(self._state_reporter.run())
+        await self._state_reconciler.run()
 
-            # Task group ensures that:
-            # 1. If one of the tasks fails then the other tasks are cancelled.
-            # 2. If Executor shuts down then all the tasks are cancelled and this function returns.
-            async with asyncio.TaskGroup() as tg:
-                tg.create_task(self._state_reporter.run())
-                tg.create_task(self._state_reconciler.run())
-        except Exception as e:
-            self._logger.error("failed running grpc mode services", exc_info=e)
-        finally:
-            # Handle task cancellation using finally.
-            if self._state_reporter is not None:
-                self._state_reporter.shutdown()
-                self._state_reporter = None
-            if self._state_reconciler is not None:
-                self._state_reconciler.shutdown()
-                self._state_reconciler = None
-
-    async def _http_mode_loop(self):
-        metric_executor_state.state("running")
-        self._startup_probe_handler.set_ready()
+    async def _http_task_runner_loop(self):
         while not self._is_shutdown:
             try:
                 async for task in self._task_fetcher.run():
@@ -341,6 +327,10 @@ class Executor:
 
     async def _shutdown(self, loop):
         self._logger.info("shutting down")
+        if self._state_reporter is not None:
+            self._state_reporter.update_executor_status(
+                ExecutorStatus.EXECUTOR_STATUS_STOPPING
+            )
         metric_executor_state.state("shutting_down")
         # There will be lots of task cancellation exceptions and "X is shutting down"
         # exceptions logged during Executor shutdown. Suppress their logs as they are
@@ -352,12 +342,13 @@ class Executor:
 
         if self._task_runner is not None:
             await self._task_runner.shutdown()
+
+        if self._channel_creator is not None:
+            await self._channel_creator.shutdown()
         if self._state_reporter is not None:
             await self._state_reporter.shutdown()
-            self._state_reporter = None
         if self._state_reconciler is not None:
             await self._state_reconciler.shutdown()
-            self._state_reconciler = None
 
         # We need to shutdown all users of FE states first,
         # otherwise states might disappear unexpectedly and we might

indexify/executor/function_executor/function_executor_state.py

@@ -1,11 +1,11 @@
 import asyncio
-from typing import Optional
-
-from indexify.task_scheduler.proto.task_scheduler_pb2 import FunctionExecutorStatus
+from typing import Any, List, Optional
 
 from .function_executor import FunctionExecutor
+from .function_executor_status import FunctionExecutorStatus, is_status_change_allowed
 from .metrics.function_executor_state import (
     metric_function_executor_state_not_locked_errors,
+    metric_function_executors_with_status,
 )
 
 
@@ -25,6 +25,7 @@ class FunctionExecutorState:
         graph_version: str,
         function_name: str,
         image_uri: Optional[str],
+        logger: Any,
     ):
         # Read only fields.
         self.id: str = id
@@ -32,69 +33,68 @@ class FunctionExecutorState:
         self.graph_name: str = graph_name
         self.function_name: str = function_name
         self.image_uri: Optional[str] = image_uri
+        self._logger: Any = logger.bind(
+            module=__name__,
+            function_executor_id=id,
+            namespace=namespace,
+            graph_name=graph_name,
+            graph_version=graph_version,
+            function_name=function_name,
+            image_uri=image_uri,
+        )
         # The lock must be held while modifying the fields below.
         self.lock: asyncio.Lock = asyncio.Lock()
+        # TODO: Move graph_version to immutable fields once we migrate to gRPC State Reconciler.
         self.graph_version: str = graph_version
-        self.is_shutdown: bool = False
-        # Set to True if a Function Executor health check ever failed.
-        self.health_check_failed: bool = False
-        # TODO: remove fields that duplicate this status field.
-        self.status: FunctionExecutorStatus = (
-            FunctionExecutorStatus.FUNCTION_EXECUTOR_STATUS_STOPPED
-        )
-        self.function_executor: Optional[FunctionExecutor] = None
-        self.running_tasks: int = 0
-        self.running_tasks_change_notifier: asyncio.Condition = asyncio.Condition(
+        self.status: FunctionExecutorStatus = FunctionExecutorStatus.DESTROYED
+        self.status_change_notifier: asyncio.Condition = asyncio.Condition(
             lock=self.lock
         )
+        self.function_executor: Optional[FunctionExecutor] = None
+        metric_function_executors_with_status.labels(status=self.status.name).inc()
 
-    def increment_running_tasks(self) -> None:
-        """Increments the number of running tasks.
-
-        The caller must hold the lock.
-        """
-        self.check_locked()
-        self.running_tasks += 1
-        self.running_tasks_change_notifier.notify_all()
-
-    def decrement_running_tasks(self) -> None:
-        """Decrements the number of running tasks.
+    async def wait_status(self, allowlist: List[FunctionExecutorStatus]) -> None:
+        """Waits until Function Executor status reaches one of the allowed values.
 
         The caller must hold the lock.
         """
         self.check_locked()
-        self.running_tasks -= 1
-        self.running_tasks_change_notifier.notify_all()
+        while self.status not in allowlist:
+            await self.status_change_notifier.wait()
 
-    async def wait_running_tasks_less(self, value: int) -> None:
-        """Waits until the number of running tasks is less than the supplied value.
+    async def set_status(self, new_status: FunctionExecutorStatus) -> None:
+        """Sets the status of the Function Executor.
 
         The caller must hold the lock.
+        Raises ValueError if the status change is not allowed.
         """
         self.check_locked()
-        while self.running_tasks >= value:
-            await self.running_tasks_change_notifier.wait()
+        if is_status_change_allowed(self.status, new_status):
+            self._logger.info(
+                "function executor status changed",
+                old_status=self.status.name,
+                new_status=new_status.name,
+            )
+            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_change_notifier.notify_all()
+        else:
+            raise ValueError(
+                f"Invalid status change from {self.status} to {new_status}"
+            )
 
     async def destroy_function_executor(self) -> None:
         """Destroys the Function Executor if it exists.
 
-        The caller must hold the lock."""
+        The caller must hold the lock.
+        """
         self.check_locked()
+        await self.set_status(FunctionExecutorStatus.DESTROYING)
         if self.function_executor is not None:
             await self.function_executor.destroy()
             self.function_executor = None
-
-    async def shutdown(self) -> None:
-        """Shuts down the state.
-
-        Called only during Executor shutdown so it's okay to fail all running and pending
-        Function Executor tasks. The state is not valid anymore after this call.
-        The caller must hold the lock.
-        """
-        self.check_locked()
-        # Pending tasks will not create a new Function Executor and won't run.
-        self.is_shutdown = True
-        await self.destroy_function_executor()
+        await self.set_status(FunctionExecutorStatus.DESTROYED)
 
     def check_locked(self) -> None:
         """Raises an exception if the lock is not held."""

indexify/executor/function_executor/function_executor_states_container.py

@@ -1,7 +1,8 @@
 import asyncio
-from typing import AsyncGenerator, Dict, Optional
+from typing import Any, AsyncGenerator, Dict, Optional
 
 from .function_executor_state import FunctionExecutorState
+from .function_executor_status import FunctionExecutorStatus
 from .metrics.function_executor_state_container import (
     metric_function_executor_states_count,
 )
@@ -10,11 +11,12 @@ from .metrics.function_executor_state_container import (
 class FunctionExecutorStatesContainer:
     """An asyncio concurrent container for the function executor states."""
 
-    def __init__(self):
+    def __init__(self, logger: Any):
         # The fields below are protected by the lock.
         self._lock: asyncio.Lock = asyncio.Lock()
         self._states: Dict[str, FunctionExecutorState] = {}
         self._is_shutdown: bool = False
+        self._logger: Any = logger.bind(module=__name__)
 
     async def get_or_create_state(
         self,
@@ -43,6 +45,7 @@ class FunctionExecutorStatesContainer:
                     graph_version=graph_version,
                     function_name=function_name,
                     image_uri=image_uri,
+                    logger=self._logger,
                 )
                 self._states[id] = state
                 metric_function_executor_states_count.set(len(self._states))
@@ -72,5 +75,8 @@ class FunctionExecutorStatesContainer:
                 # Only ongoing tasks who have a reference to the state already can see it.
                 # The state is unlocked while a task is running inside Function Executor.
                 async with state.lock:
-                    await state.shutdown()
-                    # The task running inside the Function Executor will fail because it's destroyed.
+                    await state.set_status(FunctionExecutorStatus.SHUTDOWN)
+                    if state.function_executor is not None:
+                        await state.function_executor.destroy()
+                        state.function_executor = None
+                        # The task running inside the Function Executor will fail because it's destroyed.