indexify 0.2.40__py3-none-any.whl → 0.2.41__py3-none-any.whl

indexify/cli.py

@@ -1,3 +1,8 @@
+from .logging import configure_logging_early, configure_production_logging
+
+configure_logging_early()
+
+
 import asyncio
 import os
 import shutil
@@ -18,12 +23,18 @@ from rich.text import Text
 from rich.theme import Theme
 
 from indexify.executor.agent import ExtractorAgent
-from indexify.executor.function_worker import FunctionWorker
+from indexify.function_executor.function_executor_service import (
+    FunctionExecutorService,
+)
+from indexify.function_executor.server import Server as FunctionExecutorServer
 from indexify.functions_sdk.image import (
-    DEFAULT_IMAGE_3_10,
-    DEFAULT_IMAGE_3_11,
+    LOCAL_PYTHON_VERSION,
+    GetDefaultPythonImage,
     Image,
 )
+from indexify.http_client import IndexifyClient
+
+logger = structlog.get_logger(module=__name__)
 
 custom_theme = Theme(
     {
@@ -34,10 +45,12 @@ custom_theme = Theme(
     }
 )
 
-logging = structlog.get_logger(module=__name__)
 console = Console(theme=custom_theme)
 
 app = typer.Typer(pretty_exceptions_enable=False, no_args_is_help=True)
+config_path_option: Optional[str] = typer.Option(
+    None, help="Path to the TLS configuration file"
+)
 
 
 @app.command(
@@ -149,12 +162,21 @@ def build_image(
 
 
 @app.command(help="Build default image for indexify")
-def build_default_image():
-    _build_image(image=DEFAULT_IMAGE_3_10)
-    _build_image(image=DEFAULT_IMAGE_3_11)
+def build_default_image(
+    python_version: Optional[str] = typer.Option(
+        f"{sys.version_info.major}.{sys.version_info.minor}",
+        help="Python version to use in the base image",
+    )
+):
+    image = GetDefaultPythonImage(python_version)
+
+    _build_image(image=image)
 
     console.print(
-        Text(f"Built default indexify image", style="cyan"),
+        Text(f"Built default indexify image with hash {image.hash()}\n", style="cyan"),
+        Text(
+            f"Don't forget to update your executors to run this image!", style="yellow"
+        ),
     )
 
 
@@ -164,42 +186,32 @@ def executor(
     dev: Annotated[
         bool, typer.Option("--dev", "-d", help="Run the executor in development mode")
     ] = False,
-    workers: Annotated[
-        int, typer.Option(help="number of worker processes for extraction")
-    ] = 1,
-    config_path: Optional[str] = typer.Option(
-        None, help="Path to the TLS configuration file"
-    ),
+    config_path: Optional[str] = config_path_option,
     executor_cache: Optional[str] = typer.Option(
         "~/.indexify/executor_cache", help="Path to the executor cache directory"
     ),
     name_alias: Optional[str] = typer.Option(
-        None, help="Name alias for the executor if it's spun up with the base image"
+        None, help="Image name override for the executor"
     ),
-    image_version: Optional[int] = typer.Option(
-        "1", help="Requested Image Version for this executor"
+    image_hash: Optional[str] = typer.Option(
+        None, help="Image hash override for the executor"
     ),
 ):
-    # configure structured logging
     if not dev:
-        processors = [
-            structlog.processors.dict_tracebacks,
-            structlog.processors.JSONRenderer(),
-        ]
-        structlog.configure(processors=processors)
+        configure_production_logging()
 
     id = nanoid.generate()
     executor_version = version("indexify")
-    logging.info(
+    logger.info(
         "executor started",
-        workers=workers,
         server_addr=server_addr,
         config_path=config_path,
         executor_id=id,
         executor_version=executor_version,
         executor_cache=executor_cache,
         name_alias=name_alias,
-        image_version=image_version,
+        image_hash=image_hash,
+        dev_mode=dev,
     )
 
     from pathlib import Path
@@ -211,18 +223,47 @@ def executor(
 
     agent = ExtractorAgent(
         id,
-        num_workers=workers,
         server_addr=server_addr,
         config_path=config_path,
         code_path=executor_cache,
         name_alias=name_alias,
-        image_version=image_version,
+        image_hash=image_hash,
+        development_mode=dev,
     )
 
     try:
         asyncio.get_event_loop().run_until_complete(agent.run())
     except asyncio.CancelledError:
-        logging.info("graceful shutdown")
+        logger.info("graceful shutdown")
+
+
+@app.command(help="Runs a Function Executor server")
+def function_executor(
+    function_executor_server_address: str = typer.Option(
+        help="Function Executor server address"
+    ),
+    indexify_server_address: str = typer.Option(help="Indexify server address"),
+    dev: Annotated[
+        bool, typer.Option("--dev", "-d", help="Run the executor in development mode")
+    ] = False,
+    config_path: Optional[str] = config_path_option,
+):
+    if not dev:
+        configure_production_logging()
+
+    logger.info(
+        "starting function executor server",
+        function_executor_server_address=function_executor_server_address,
+        indexify_server_address=indexify_server_address,
+        config_path=config_path,
+    )
+
+    FunctionExecutorServer(
+        server_address=function_executor_server_address,
+        service=FunctionExecutorService(
+            indexify_server_address=indexify_server_address, config_path=config_path
+        ),
+    ).run()
 
 
 def _create_image(image: Image, python_sdk_path):
@@ -234,6 +275,7 @@ def _create_image(image: Image, python_sdk_path):
 
 
 def _build_image(image: Image, python_sdk_path: Optional[str] = None):
+
     try:
         import docker
 
@@ -246,24 +288,31 @@ def _build_image(image: Image, python_sdk_path: Optional[str] = None):
         )
         exit(-1)
 
-    docker_file = f"""
-FROM {image._base_image}
-
-RUN mkdir -p ~/.indexify
+    docker_contents = [
+        f"FROM {image._base_image}",
+        "RUN mkdir -p ~/.indexify",
+        "RUN touch ~/.indexify/image_name",
+        f"RUN echo {image._image_name} > ~/.indexify/image_name",
+        f"RUN echo {image.hash()} > ~/.indexify/image_hash",
+        "WORKDIR /app",
+    ]
 
-RUN touch ~/.indexify/image_name
+    docker_contents.extend(["RUN " + i for i in image._run_strs])
 
-RUN  echo {image._image_name} > ~/.indexify/image_name
-
-WORKDIR /app
-
-"""
+    if python_sdk_path is not None:
+        logging.info(
+            f"Building image {image._image_name} with local version of the SDK"
+        )
+        if not os.path.exists(python_sdk_path):
+            print(f"error: {python_sdk_path} does not exist")
+            os.exit(1)
+        docker_contents.append(f"COPY {python_sdk_path} /app/python-sdk")
+        docker_contents.append("RUN (cd /app/python-sdk && pip install .)")
+    else:
+        docker_contents.append(f"RUN pip install indexify=={image._sdk_version}")
 
-    run_strs = ["RUN " + i for i in image._run_strs]
+    docker_file = "\n".join(docker_contents)
 
-    docker_file += "\n".join(run_strs)
-    print(os.getcwd())
-    import docker
     import docker.api.build
 
     docker.api.build.process_dockerfile = lambda dockerfile, path: (
@@ -271,15 +320,6 @@ WORKDIR /app
         dockerfile,
     )
 
-    if python_sdk_path is not None:
-        if not os.path.exists(python_sdk_path):
-            print(f"error: {python_sdk_path} does not exist")
-            os.exit(1)
-        docker_file += f"\nCOPY {python_sdk_path} /app/python-sdk"
-        docker_file += f"\nRUN (cd /app/python-sdk && pip install .)"
-    else:
-        docker_file += f"\nRUN pip install indexify"
-
     console.print("Creating image using Dockerfile contents:", style="cyan bold")
     print(f"{docker_file}")
 

indexify/executor/agent.py

@@ -1,109 +1,87 @@
 import asyncio
-import json
-from concurrent.futures.process import BrokenProcessPool
-from importlib.metadata import version
 from pathlib import Path
 from typing import Dict, List, Optional
 
 import structlog
-from httpx_sse import aconnect_sse
-from pydantic import BaseModel
 
-from indexify.common_util import get_httpx_client
-from indexify.functions_sdk.data_objects import (
+from .downloader import Downloader
+from .executor_tasks import DownloadGraphTask, DownloadInputsTask, RunTask
+from .function_executor.process_function_executor_factory import (
+    ProcessFunctionExecutorFactory,
+)
+from .function_worker import (
+    FunctionWorker,
+    FunctionWorkerInput,
     FunctionWorkerOutput,
-    IndexifyData,
 )
-from indexify.http_client import IndexifyClient
-
-from .api_objects import ExecutorMetadata, Task
-from .downloader import DownloadedInputs, Downloader
-from .executor_tasks import DownloadGraphTask, DownloadInputTask, ExtractTask
-from .function_worker import FunctionWorker
-from .runtime_probes import ProbeInfo, RuntimeProbes
+from .task_fetcher import TaskFetcher
 from .task_reporter import TaskReporter
 from .task_store import CompletedTask, TaskStore
 
-logging = structlog.get_logger(module=__name__)
-
-
-class FunctionInput(BaseModel):
-    task_id: str
-    namespace: str
-    compute_graph: str
-    function: str
-    input: IndexifyData
-    init_value: Optional[IndexifyData] = None
+logger = structlog.get_logger(module=__name__)
 
 
 class ExtractorAgent:
     def __init__(
         self,
         executor_id: str,
-        num_workers,
         code_path: Path,
         server_addr: str = "localhost:8900",
+        development_mode: bool = False,
         config_path: Optional[str] = None,
         name_alias: Optional[str] = None,
-        image_version: Optional[int] = None,
+        image_hash: Optional[str] = None,
     ):
-        self.name_alias = name_alias
-        self.image_version = image_version
         self._config_path = config_path
-        self._probe = RuntimeProbes()
-
-        self.num_workers = num_workers
+        protocol: str = "http"
         if config_path:
-            logging.info("running the extractor with TLS enabled")
-            self._protocol = "https"
-        else:
-            self._protocol = "http"
+            logger.info("running the extractor with TLS enabled")
+            protocol = "https"
 
         self._task_store: TaskStore = TaskStore()
-        self._executor_id = executor_id
         self._function_worker = FunctionWorker(
-            workers=num_workers,
-            indexify_client=IndexifyClient(
-                service_url=f"{self._protocol}://{server_addr}",
+            function_executor_factory=ProcessFunctionExecutorFactory(
+                indexify_server_address=server_addr,
+                development_mode=development_mode,
                 config_path=config_path,
-            ),
+            )
         )
         self._has_registered = False
         self._server_addr = server_addr
-        self._base_url = f"{self._protocol}://{self._server_addr}"
+        self._base_url = f"{protocol}://{self._server_addr}"
         self._code_path = code_path
         self._downloader = Downloader(
-            code_path=code_path, base_url=self._base_url, config_path=self._config_path
+            code_path=code_path, base_url=self._base_url, config_path=config_path
+        )
+        self._task_fetcher = TaskFetcher(
+            protocol=protocol,
+            indexify_server_addr=self._server_addr,
+            executor_id=executor_id,
+            name_alias=name_alias,
+            image_hash=image_hash,
+            config_path=config_path,
         )
-        self._max_queued_tasks = 10
         self._task_reporter = TaskReporter(
             base_url=self._base_url,
-            executor_id=self._executor_id,
+            executor_id=executor_id,
             config_path=self._config_path,
         )
 
     async def task_completion_reporter(self):
-        logging.info("starting task completion reporter")
+        logger.info("starting task completion reporter")
         # We should copy only the keys and not the values
         while True:
             outcomes = await self._task_store.task_outcomes()
             for task_outcome in outcomes:
-                retryStr = (
-                    f"\nRetries: {task_outcome.reporting_retries}"
-                    if task_outcome.reporting_retries > 0
-                    else ""
-                )
-                outcome = task_outcome.task_outcome
-                style_outcome = (
-                    f"[bold red] {outcome} [/]"
-                    if "fail" in outcome
-                    else f"[bold green] {outcome} [/]"
-                )
-                logging.info(
+                logger.info(
                     "reporting_task_outcome",
                     task_id=task_outcome.task.id,
                     fn_name=task_outcome.task.compute_fn,
-                    num_outputs=len(task_outcome.outputs or []),
+                    num_outputs=(
+                        len(task_outcome.function_output.outputs)
+                        if task_outcome.function_output is not None
+                        else 0
+                    ),
                     router_output=task_outcome.router_output,
                     outcome=task_outcome.task_outcome,
                     retries=task_outcome.reporting_retries,
@@ -114,10 +92,10 @@ class ExtractorAgent:
                     self._task_reporter.report_task_outcome(completed_task=task_outcome)
                 except Exception as e:
                     # The connection was dropped in the middle of the reporting, process, retry
-                    logging.error(
+                    logger.error(
                         "failed_to_report_task",
                         task_id=task_outcome.task.id,
-                        exception=f"exception: {type(e).__name__}({e})",
+                        exc_info=e,
                         retries=task_outcome.reporting_retries,
                     )
                     task_outcome.reporting_retries += 1
@@ -127,30 +105,13 @@ class ExtractorAgent:
                 self._task_store.mark_reported(task_id=task_outcome.task.id)
 
     async def task_launcher(self):
-        async_tasks: List[asyncio.Task] = []
-        fn_queue: List[FunctionInput] = []
-
-        async_tasks.append(
+        async_tasks: List[asyncio.Task] = [
             asyncio.create_task(
                 self._task_store.get_runnable_tasks(), name="get_runnable_tasks"
             )
-        )
+        ]
 
         while True:
-            fn: FunctionInput
-            for fn in fn_queue:
-                task: Task = self._task_store.get_task(fn.task_id)
-                async_tasks.append(
-                    ExtractTask(
-                        function_worker=self._function_worker,
-                        task=task,
-                        input=fn.input,
-                        code_path=f"{self._code_path}/{task.namespace}/{task.compute_graph}.{task.graph_version}",
-                        init_value=fn.init_value,
-                    )
-                )
-
-            fn_queue = []
             done, pending = await asyncio.wait(
                 async_tasks, return_when=asyncio.FIRST_COMPLETED
             )
@@ -159,16 +120,19 @@ class ExtractorAgent:
             for async_task in done:
                 if async_task.get_name() == "get_runnable_tasks":
                     if async_task.exception():
-                        logging.error(
+                        logger.error(
                             "task_launcher_error, failed to get runnable tasks",
-                            exception=async_task.exception(),
+                            exc_info=async_task.exception(),
                         )
                         continue
                     result: Dict[str, Task] = await async_task
                     task: Task
                     for _, task in result.items():
                         async_tasks.append(
-                            DownloadGraphTask(task=task, downloader=self._downloader)
+                            DownloadGraphTask(
+                                function_worker_input=FunctionWorkerInput(task=task),
+                                downloader=self._downloader,
+                            )
                         )
                     async_tasks.append(
                         asyncio.create_task(
@@ -178,58 +142,60 @@ class ExtractorAgent:
                     )
                 elif async_task.get_name() == "download_graph":
                     if async_task.exception():
-                        logging.error(
+                        logger.error(
                             "task_launcher_error, failed to download graph",
-                            exception=async_task.exception(),
+                            exc_info=async_task.exception(),
                         )
                         completed_task = CompletedTask(
-                            task=async_task.task,
-                            outputs=[],
+                            task=async_task.function_worker_input.task,
                             task_outcome="failure",
                         )
                         self._task_store.complete(outcome=completed_task)
                         continue
+                    async_task: DownloadGraphTask
+                    function_worker_input: FunctionWorkerInput = (
+                        async_task.function_worker_input
+                    )
+                    function_worker_input.graph = await async_task
                     async_tasks.append(
-                        DownloadInputTask(
-                            task=async_task.task, downloader=self._downloader
+                        DownloadInputsTask(
+                            function_worker_input=function_worker_input,
+                            downloader=self._downloader,
                         )
                     )
-                elif async_task.get_name() == "download_input":
+                elif async_task.get_name() == "download_inputs":
                     if async_task.exception():
-                        logging.error(
-                            "task_launcher_error, failed to download input",
-                            exception=str(async_task.exception()),
+                        logger.error(
+                            "task_launcher_error, failed to download inputs",
+                            exc_info=async_task.exception(),
                         )
                         completed_task = CompletedTask(
-                            task=async_task.task,
-                            outputs=[],
+                            task=async_task.function_worker_input.task,
                             task_outcome="failure",
                         )
                         self._task_store.complete(outcome=completed_task)
                         continue
-                    downloaded_inputs: DownloadedInputs = await async_task
-                    task: Task = async_task.task
-                    fn_queue.append(
-                        FunctionInput(
-                            task_id=task.id,
-                            namespace=task.namespace,
-                            compute_graph=task.compute_graph,
-                            function=task.compute_fn,
-                            input=downloaded_inputs.input,
-                            init_value=downloaded_inputs.init_value,
+                    async_task: DownloadInputsTask
+                    function_worker_input: FunctionWorkerInput = (
+                        async_task.function_worker_input
+                    )
+                    function_worker_input.function_input = await async_task
+                    async_tasks.append(
+                        RunTask(
+                            function_worker=self._function_worker,
+                            function_worker_input=function_worker_input,
                         )
                     )
-                elif async_task.get_name() == "run_function":
+                elif async_task.get_name() == "run_task":
                     if async_task.exception():
                         completed_task = CompletedTask(
-                            task=async_task.task,
+                            task=async_task.function_worker_input.task,
                             task_outcome="failure",
-                            outputs=[],
                             stderr=str(async_task.exception()),
                         )
                         self._task_store.complete(outcome=completed_task)
                         continue
-                    async_task: ExtractTask
+                    async_task: RunTask
                     try:
                         outputs: FunctionWorkerOutput = await async_task
                         if not outputs.success:
@@ -238,113 +204,59 @@ class ExtractorAgent:
                             task_outcome = "success"
 
                         completed_task = CompletedTask(
-                            task=async_task.task,
+                            task=async_task.function_worker_input.task,
                             task_outcome=task_outcome,
-                            outputs=outputs.fn_outputs,
+                            function_output=outputs.function_output,
                             router_output=outputs.router_output,
                             stdout=outputs.stdout,
                             stderr=outputs.stderr,
                             reducer=outputs.reducer,
                         )
                         self._task_store.complete(outcome=completed_task)
-                    except BrokenProcessPool:
-                        self._task_store.retriable_failure(async_task.task.id)
-                        continue
                     except Exception as e:
-                        logging.error(
+                        logger.error(
                             "failed to execute task",
-                            task_id=async_task.task.id,
-                            exception=str(e),
+                            task_id=async_task.function_worker_input.task.id,
+                            exc_info=e,
                         )
                         completed_task = CompletedTask(
-                            task=async_task.task,
+                            task=async_task.function_worker_input.task,
                             task_outcome="failure",
-                            outputs=[],
                         )
                         self._task_store.complete(outcome=completed_task)
                         continue
 
+    async def _main_loop(self):
+        """Fetches incoming tasks from the server and starts their processing."""
+        self._should_run = True
+        while self._should_run:
+            try:
+                async for task in self._task_fetcher.run():
+                    self._task_store.add_tasks([task])
+            except Exception as e:
+                logger.error("failed fetching tasks, retrying in 5 seconds", exc_info=e)
+                await asyncio.sleep(5)
+                continue
+
     async def run(self):
         import signal
 
         asyncio.get_event_loop().add_signal_handler(
             signal.SIGINT, self.shutdown, asyncio.get_event_loop()
         )
+        asyncio.get_event_loop().add_signal_handler(
+            signal.SIGTERM, self.shutdown, asyncio.get_event_loop()
+        )
         asyncio.create_task(self.task_launcher())
         asyncio.create_task(self.task_completion_reporter())
-        self._should_run = True
-        while self._should_run:
-            url = f"{self._protocol}://{self._server_addr}/internal/executors/{self._executor_id}/tasks"
-            runtime_probe: ProbeInfo = self._probe.probe()
-
-            executor_version = version("indexify")
-
-            image_name = (
-                self.name_alias
-                if self.name_alias is not None
-                else runtime_probe.image_name
-            )
-
-            image_version: int = (
-                self.image_version
-                if self.image_version is not None
-                else runtime_probe.image_version
-            )
-
-            data = ExecutorMetadata(
-                id=self._executor_id,
-                executor_version=executor_version,
-                addr="",
-                image_name=image_name,
-                image_version=image_version,
-                labels=runtime_probe.labels,
-            ).model_dump()
-            logging.info(
-                "registering_executor",
-                executor_id=self._executor_id,
-                url=url,
-                executor_version=executor_version,
-            )
-            try:
-                async with get_httpx_client(self._config_path, True) as client:
-                    async with aconnect_sse(
-                        client,
-                        "POST",
-                        url,
-                        json=data,
-                        headers={"Content-Type": "application/json"},
-                    ) as event_source:
-                        if not event_source.response.is_success:
-                            resp = await event_source.response.aread()
-                            logging.error(
-                                f"failed to register",
-                                resp=str(resp),
-                                status_code=event_source.response.status_code,
-                            )
-                            await asyncio.sleep(5)
-                            continue
-                        logging.info(
-                            "executor_registered", executor_id=self._executor_id
-                        )
-                        async for sse in event_source.aiter_sse():
-                            data = json.loads(sse.data)
-                            tasks = []
-                            for task_dict in data:
-                                tasks.append(
-                                    Task.model_validate(task_dict, strict=False)
-                                )
-                            self._task_store.add_tasks(tasks)
-            except Exception as e:
-                logging.error(f"failed to register: {e}")
-                await asyncio.sleep(5)
-                continue
+        await self._main_loop()
 
     async def _shutdown(self, loop):
-        logging.info("shutting_down")
+        logger.info("shutting_down")
         self._should_run = False
+        await self._function_worker.shutdown()
         for task in asyncio.all_tasks(loop):
             task.cancel()
 
     def shutdown(self, loop):
-        self._function_worker.shutdown()
         loop.create_task(self._shutdown(loop))