indexify 0.0.43__py3-none-any.whl → 0.2__py3-none-any.whl

indexify/error.py

@@ -1,30 +1,3 @@
-class Error(Exception):
-    status: str
-    message: str
-
-    def __init__(self, status: str, message: str):
-        self.status = status
-        self.message = message
-
-    @staticmethod
-    def from_tonic_error_string(url: str, error: str) -> "Error":
-        data = error.split(", ")
-
-        message = data[1].split(": ", 1)[1]
-        if message.startswith('"') and message.endswith('"'):
-            message = message[1:-1]
-
-        status = "GeneralError"
-        if "extraction_graph" in url:
-            status = "ExtractionGraphError"
-        elif "search" in url:
-            status = "SearchError"
-
-        error = Error(status, message)
-        return error
-
-    def __str__(self):
-        return f"{self.status} | {self.message.capitalize()}"
-
-    def __repr__(self):
-        return f"Error(status={self.status!r}, message={self.message!r})"
+class ApiException(Exception):
+    def __init__(self, message: str) -> None:
+        super().__init__(message)

indexify/executor/agent.py

@@ -0,0 +1,362 @@
+import asyncio
+import json
+import ssl
+import traceback
+from concurrent.futures.process import BrokenProcessPool
+from typing import Dict, List, Optional
+
+import httpx
+import yaml
+from httpx_sse import aconnect_sse
+from pydantic import BaseModel
+from rich.console import Console
+from rich.panel import Panel
+from rich.text import Text
+from rich.theme import Theme
+
+from indexify.functions_sdk.data_objects import (
+    FunctionWorkerOutput,
+    IndexifyData,
+    RouterOutput,
+)
+
+from .api_objects import ExecutorMetadata, Task
+from .downloader import DownloadedInputs, Downloader
+from .executor_tasks import DownloadGraphTask, DownloadInputTask, ExtractTask
+from .function_worker import FunctionWorker
+from .task_reporter import TaskReporter
+from .task_store import CompletedTask, TaskStore
+
+custom_theme = Theme(
+    {
+        "info": "cyan",
+        "warning": "yellow",
+        "error": "red",
+        "success": "green",
+    }
+)
+
+console = Console(theme=custom_theme)
+
+
+class FunctionInput(BaseModel):
+    task_id: str
+    namespace: str
+    compute_graph: str
+    function: str
+    input: IndexifyData
+    init_value: Optional[IndexifyData] = None
+
+
+class ExtractorAgent:
+    def __init__(
+        self,
+        executor_id: str,
+        num_workers,
+        code_path: str,
+        function_worker: FunctionWorker,
+        server_addr: str = "localhost:8900",
+        config_path: Optional[str] = None,
+    ):
+        self.num_workers = num_workers
+        self._use_tls = False
+        if config_path:
+            with open(config_path, "r") as f:
+                config = yaml.safe_load(f)
+                self._config = config
+            if config.get("use_tls", False):
+                console.print(
+                    "Running the extractor with TLS enabled", style="cyan bold"
+                )
+                self._use_tls = True
+                tls_config = config["tls_config"]
+                self._ssl_context = ssl.create_default_context(
+                    ssl.Purpose.SERVER_AUTH, cafile=tls_config["ca_bundle_path"]
+                )
+                self._ssl_context.load_cert_chain(
+                    certfile=tls_config["cert_path"], keyfile=tls_config["key_path"]
+                )
+                self._protocol = "wss"
+                self._tls_config = tls_config
+            else:
+                self._ssl_context = None
+                self._protocol = "ws"
+        else:
+            self._ssl_context = None
+            self._protocol = "http"
+            self._config = {}
+
+        self._task_store: TaskStore = TaskStore()
+        self._executor_id = executor_id
+        self._function_worker = function_worker
+        self._has_registered = False
+        self._server_addr = server_addr
+        self._base_url = f"{self._protocol}://{self._server_addr}"
+        self._code_path = code_path
+        self._downloader = Downloader(code_path=code_path, base_url=self._base_url)
+        self._max_queued_tasks = 10
+        self._task_reporter = TaskReporter(
+            base_url=self._base_url, executor_id=self._executor_id
+        )
+
+    async def task_completion_reporter(self):
+        console.print(Text("Starting task completion reporter", style="bold cyan"))
+        # We should copy only the keys and not the values
+        url = f"{self._protocol}://{self._server_addr}/write_content"
+        while True:
+            outcomes = await self._task_store.task_outcomes()
+            for task_outcome in outcomes:
+                outcome = task_outcome.task_outcome
+                style_outcome = (
+                    f"[bold red] {outcome} [/]"
+                    if "fail" in outcome
+                    else f"[bold green] {outcome} [/]"
+                )
+                console.print(
+                    Panel(
+                        f"Reporting outcome of task {task_outcome.task.id}\n"
+                        f"Outcome: {style_outcome}\n"
+                        f"Outputs: {len(task_outcome.outputs or [])}  Router Output: {task_outcome.router_output}",
+                        title="Task Completion",
+                        border_style="info",
+                    )
+                )
+
+                try:
+                    # Send task outcome to the server
+                    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
+                    console.print(
+                        Panel(
+                            f"Failed to report task {task_outcome.task.id}\n"
+                            f"Exception: {e}\nRetrying...",
+                            title="Reporting Error",
+                            border_style="error",
+                        )
+                    )
+                    await asyncio.sleep(5)
+                    continue
+
+                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(
+            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
+            )
+
+            async_tasks: List[asyncio.Task] = list(pending)
+            for async_task in done:
+                if async_task.get_name() == "get_runnable_tasks":
+                    if async_task.exception():
+                        console.print(
+                            Text("Task Launcher Error: ", style="red bold")
+                            + Text(
+                                f"Failed to get runnable tasks: {async_task.exception()}",
+                                style="red",
+                            )
+                        )
+                        continue
+                    result: Dict[str, Task] = await async_task
+                    task: Task
+                    for _, task in result.items():
+                        async_tasks.append(
+                            DownloadGraphTask(task=task, downloader=self._downloader)
+                        )
+                    async_tasks.append(
+                        asyncio.create_task(
+                            self._task_store.get_runnable_tasks(),
+                            name="get_runnable_tasks",
+                        )
+                    )
+                elif async_task.get_name() == "download_graph":
+                    if async_task.exception():
+                        console.print(
+                            Text(
+                                f"Failed to download graph for task {async_task.task.id}\n",
+                                style="red bold",
+                            )
+                            + Text(f"Exception: {async_task.exception()}", style="red")
+                        )
+                        completed_task = CompletedTask(
+                            task=async_task.task,
+                            outputs=[],
+                            task_outcome="failure",
+                        )
+                        self._task_store.complete(outcome=completed_task)
+                        continue
+                    async_tasks.append(
+                        DownloadInputTask(
+                            task=async_task.task, downloader=self._downloader
+                        )
+                    )
+                elif async_task.get_name() == "download_input":
+                    if async_task.exception():
+                        console.print(
+                            Text(
+                                f"Failed to download input for task {async_task.task.id}\n",
+                                style="red bold",
+                            )
+                            + Text(f"Exception: {async_task.exception()}", style="red")
+                        )
+                        completed_task = CompletedTask(
+                            task=async_task.task,
+                            outputs=[],
+                            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,
+                        )
+                    )
+                elif async_task.get_name() == "run_function":
+                    if async_task.exception():
+                        completed_task = CompletedTask(
+                            task=async_task.task,
+                            task_outcome="failure",
+                            outputs=[],
+                            errors=str(async_task.exception()),
+                        )
+                        self._task_store.complete(outcome=completed_task)
+                        continue
+                    async_task: ExtractTask
+                    try:
+                        outputs: FunctionWorkerOutput = await async_task
+                        if not outputs.success:
+                            task_outcome = "failure"
+                        else:
+                            task_outcome = "success"
+
+                        completed_task = CompletedTask(
+                            task=async_task.task,
+                            task_outcome=task_outcome,
+                            outputs=outputs.fn_outputs,
+                            router_output=outputs.router_output,
+                            errors=outputs.exception,
+                            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:
+                        console.print(
+                            Text(
+                                f"Failed to execute task {async_task.task.id}\n",
+                                style="red bold",
+                            )
+                            + Text(f"Exception: {e}", style="red")
+                        )
+                        completed_task = CompletedTask(
+                            task=async_task.task,
+                            task_outcome="failure",
+                            outputs=[],
+                        )
+                        self._task_store.complete(outcome=completed_task)
+                        continue
+
+    async def run(self):
+        import signal
+
+        asyncio.get_event_loop().add_signal_handler(
+            signal.SIGINT, 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:
+            self._protocol = "http"
+            url = f"{self._protocol}://{self._server_addr}/internal/executors/{self._executor_id}/tasks"
+
+            def to_sentence_case(snake_str):
+                words = snake_str.split("_")
+                return words[0].capitalize() + "" + " ".join(words[1:])
+
+            data = ExecutorMetadata(
+                id=self._executor_id,
+                address="",
+                runner_name="extractor",
+                labels={},
+            ).model_dump()
+
+            panel_content = "\n".join(
+                [f"{to_sentence_case(key)}: {value}" for key, value in data.items()]
+            )
+            console.print(
+                Panel(
+                    panel_content,
+                    title="Attempting to Register Executor",
+                    border_style="cyan",
+                )
+            )
+
+            try:
+                async with httpx.AsyncClient() as client:
+                    async with aconnect_sse(
+                        client,
+                        "POST",
+                        url,
+                        json=data,
+                        headers={"Content-Type": "application/json"},
+                    ) as event_source:
+                        console.print(
+                            Text("Executor registered successfully", style="bold green")
+                        )
+                        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:
+                console.print(
+                    Text("Registration Error: ", style="red bold")
+                    + Text(f"Failed to register: {e}", style="red")
+                )
+                await asyncio.sleep(5)
+                continue
+
+    async def _shutdown(self, loop):
+        console.print(Text("Shutting down agent...", style="bold yellow"))
+        self._should_run = False
+        for task in asyncio.all_tasks(loop):
+            task.cancel()
+
+    def shutdown(self, loop):
+        self._function_worker.shutdown()
+        loop.create_task(self._shutdown(loop))

indexify/executor/api_objects.py

@@ -0,0 +1,43 @@
+from typing import Any, Dict, List, Optional
+
+from pydantic import BaseModel, Json
+
+from indexify.functions_sdk.data_objects import IndexifyData
+
+
+class Task(BaseModel):
+    id: str
+    namespace: str
+    compute_graph: str
+    compute_fn: str
+    invocation_id: str
+    input_key: str
+    reducer_output_id: Optional[str] = None
+    graph_version: int
+
+
+class ExecutorMetadata(BaseModel):
+    id: str
+    address: str
+    runner_name: str
+    labels: Dict[str, Any]
+
+
+class RouterOutput(BaseModel):
+    edges: List[str]
+
+
+class FnOutput(BaseModel):
+    payload: Json
+
+
+class TaskResult(BaseModel):
+    router_output: Optional[RouterOutput] = None
+    outcome: str
+    namespace: str
+    compute_graph: str
+    compute_fn: str
+    invocation_id: str
+    executor_id: str
+    task_id: str
+    reducer: bool = False

indexify/executor/downloader.py

@@ -0,0 +1,124 @@
+import os
+from typing import Optional
+
+import httpx
+from pydantic import BaseModel
+from rich.console import Console
+from rich.panel import Panel
+from rich.theme import Theme
+
+from indexify.functions_sdk.data_objects import IndexifyData
+from indexify.functions_sdk.object_serializer import MsgPackSerializer
+
+from .api_objects import Task
+
+custom_theme = Theme(
+    {
+        "info": "cyan",
+        "warning": "yellow",
+        "error": "red",
+    }
+)
+
+console = Console(theme=custom_theme)
+
+
+class DownloadedInputs(BaseModel):
+    input: IndexifyData
+    init_value: Optional[IndexifyData] = None
+
+
+class Downloader:
+    def __init__(self, code_path: str, base_url: str):
+        self.code_path = code_path
+        self.base_url = base_url
+
+    async def download_graph(self, namespace: str, name: str, version: int) -> str:
+        path = os.path.join(self.code_path, namespace, f"{name}.{version}")
+        if os.path.exists(path):
+            return path
+
+        console.print(
+            Panel(
+                f"Downloading graph: {name}\nPath: {path}",
+                title="downloader",
+                border_style="cyan",
+            )
+        )
+
+        response = httpx.get(
+            f"{self.base_url}/internal/namespaces/{namespace}/compute_graphs/{name}/code"
+        )
+        try:
+            response.raise_for_status()
+        except httpx.HTTPStatusError as e:
+            console.print(
+                Panel(
+                    f"Failed to download graph: {name}\nError: {response.text}",
+                    title="downloader error",
+                    border_style="error",
+                )
+            )
+            raise
+
+        os.makedirs(os.path.dirname(path), exist_ok=True)
+        with open(path, "wb") as f:
+            f.write(response.content)
+        return path
+
+    async def download_input(self, task: Task) -> IndexifyData:
+        input_id = task.input_key.split("|")[-1]
+        if task.invocation_id == input_id:
+            url = f"{self.base_url}/namespaces/{task.namespace}/compute_graphs/{task.compute_graph}/invocations/{task.invocation_id}/payload"
+        else:
+            url = f"{self.base_url}/internal/fn_outputs/{task.input_key}"
+
+        reducer_url = None
+        if task.reducer_output_id:
+            reducer_url = f"{self.base_url}/namespaces/{task.namespace}/compute_graphs/{task.compute_graph}/invocations/{task.invocation_id}/fn/{task.compute_fn}/{task.reducer_output_id}"
+
+        console.print(
+            Panel(
+                f"downloading input\nURL: {url} \n reducer input URL: {reducer_url}",
+                title="downloader",
+                border_style="cyan",
+            )
+        )
+
+        response = httpx.get(url)
+        try:
+            response.raise_for_status()
+        except httpx.HTTPStatusError as e:
+            console.print(
+                Panel(
+                    f"failed to download input: {task.input_key}\nError: {response.text}",
+                    title="downloader error",
+                    border_style="error",
+                )
+            )
+            raise
+
+        if task.invocation_id == input_id:
+            return DownloadedInputs(
+                input=IndexifyData(payload=response.content, id=input_id)
+            )
+
+        init_value = None
+        if reducer_url:
+            init_value = httpx.get(reducer_url)
+            try:
+                init_value.raise_for_status()
+            except httpx.HTTPStatusError as e:
+                console.print(
+                    Panel(
+                        f"failed to download reducer output: {task.reducer_output_id}\nError: {init_value.text}",
+                        title="downloader error",
+                        border_style="error",
+                    )
+                )
+                raise
+            init_value = MsgPackSerializer.deserialize(init_value.content)
+
+        return DownloadedInputs(
+            input=MsgPackSerializer.deserialize(response.content), init_value=init_value
+        )

indexify/executor/executor_tasks.py

@@ -0,0 +1,72 @@
+import asyncio
+from typing import Optional
+
+from indexify.functions_sdk.data_objects import IndexifyData
+
+from .api_objects import Task
+from .downloader import Downloader
+from .function_worker import FunctionWorker
+
+
+class DownloadGraphTask(asyncio.Task):
+    def __init__(
+        self,
+        *,
+        task: Task,
+        downloader: Downloader,
+        **kwargs,
+    ):
+        kwargs["name"] = "download_graph"
+        kwargs["loop"] = asyncio.get_event_loop()
+        super().__init__(
+            downloader.download_graph(
+                task.namespace, task.compute_graph, task.graph_version
+            ),
+            **kwargs,
+        )
+        self.task = task
+
+
+class DownloadInputTask(asyncio.Task):
+    def __init__(
+        self,
+        *,
+        task: Task,
+        downloader: Downloader,
+        **kwargs,
+    ):
+        kwargs["name"] = "download_input"
+        kwargs["loop"] = asyncio.get_event_loop()
+        super().__init__(
+            downloader.download_input(task),
+            **kwargs,
+        )
+        self.task = task
+
+
+class ExtractTask(asyncio.Task):
+    def __init__(
+        self,
+        *,
+        function_worker: FunctionWorker,
+        task: Task,
+        input: IndexifyData,
+        init_value: Optional[IndexifyData] = None,
+        code_path: str,
+        **kwargs,
+    ):
+        kwargs["name"] = "run_function"
+        kwargs["loop"] = asyncio.get_event_loop()
+        super().__init__(
+            function_worker.async_submit(
+                namespace=task.namespace,
+                graph_name=task.compute_graph,
+                fn_name=task.compute_fn,
+                input=input,
+                init_value=init_value,
+                code_path=code_path,
+                version=task.graph_version,
+            ),
+            **kwargs,
+        )
+        self.task = task