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

indexify/executor/function_worker.py

@@ -0,0 +1,177 @@
+import asyncio
+import traceback
+from concurrent.futures.process import BrokenProcessPool
+from typing import Dict, List, Optional, Union
+
+from pydantic import BaseModel
+from rich import print
+
+from indexify.functions_sdk.data_objects import (
+    FunctionWorkerOutput,
+    IndexifyData,
+    RouterOutput,
+)
+from indexify.functions_sdk.graph import Graph
+from indexify.functions_sdk.indexify_functions import IndexifyFunctionWrapper
+
+graphs: Dict[str, Graph] = {}
+function_wrapper_map: Dict[str, IndexifyFunctionWrapper] = {}
+
+import concurrent.futures
+import io
+from contextlib import redirect_stderr, redirect_stdout
+
+
+class FunctionRunException(Exception):
+    def __init__(
+        self, exception: Exception, stdout: str, stderr: str, is_reducer: bool
+    ):
+        super().__init__(str(exception))
+        self.exception = exception
+        self.stdout = stdout
+        self.stderr = stderr
+        self.is_reducer = is_reducer
+
+
+class FunctionOutput(BaseModel):
+    fn_outputs: Optional[List[IndexifyData]]
+    router_output: Optional[RouterOutput]
+    reducer: bool = False
+    success: bool = True
+    exception: Optional[str] = None
+    stdout: str = ""
+    stderr: str = ""
+
+
+def _load_function(
+    namespace: str, graph_name: str, fn_name: str, code_path: str, version: int
+):
+    """Load an extractor to the memory: extractor_wrapper_map."""
+    global function_wrapper_map
+    key = f"{namespace}/{graph_name}/{version}/{fn_name}"
+    if key in function_wrapper_map:
+        return
+    graph = Graph.from_path(code_path)
+    function_wrapper = graph.get_function(fn_name)
+    function_wrapper_map[key] = function_wrapper
+    graph_key = f"{namespace}/{graph_name}/{version}"
+    graphs[graph_key] = graph
+
+
+class FunctionWorker:
+    def __init__(self, workers: int = 1) -> None:
+        self._executor: concurrent.futures.ProcessPoolExecutor = (
+            concurrent.futures.ProcessPoolExecutor(max_workers=workers)
+        )
+
+    async def async_submit(
+        self,
+        namespace: str,
+        graph_name: str,
+        fn_name: str,
+        input: IndexifyData,
+        code_path: str,
+        version: int,
+        init_value: Optional[IndexifyData] = None,
+    ) -> FunctionWorkerOutput:
+        try:
+            result = await asyncio.get_running_loop().run_in_executor(
+                self._executor,
+                _run_function,
+                namespace,
+                graph_name,
+                fn_name,
+                input,
+                code_path,
+                version,
+                init_value,
+            )
+        except BrokenProcessPool as mp:
+            self._executor.shutdown(wait=True, cancel_futures=True)
+            traceback.print_exc()
+            raise mp
+        except FunctionRunException as e:
+            print(e)
+            print(traceback.format_exc())
+            return FunctionWorkerOutput(
+                exception=str(e),
+                stdout=e.stdout,
+                stderr=e.stderr,
+                reducer=e.is_reducer,
+                success=False,
+            )
+
+        return FunctionWorkerOutput(
+            fn_outputs=result.fn_outputs,
+            router_output=result.router_output,
+            exception=result.exception,
+            stdout=result.stdout,
+            stderr=result.stderr,
+            reducer=result.reducer,
+            success=result.success,
+        )
+
+    def shutdown(self):
+        self._executor.shutdown(wait=True, cancel_futures=True)
+
+
+def _run_function(
+    namespace: str,
+    graph_name: str,
+    fn_name: str,
+    input: IndexifyData,
+    code_path: str,
+    version: int,
+    init_value: Optional[IndexifyData] = None,
+) -> FunctionOutput:
+    import io
+    import traceback
+    from contextlib import redirect_stderr, redirect_stdout
+
+    stdout_capture = io.StringIO()
+    stderr_capture = io.StringIO()
+    is_reducer = False
+    router_output = None
+    fn_output = None
+    has_failed = False
+    exception_msg = None
+    with redirect_stdout(stdout_capture), redirect_stderr(stderr_capture):
+        try:
+            key = f"{namespace}/{graph_name}/{version}/{fn_name}"
+            if key not in function_wrapper_map:
+                _load_function(namespace, graph_name, fn_name, code_path, version)
+
+            graph: Graph = graphs[f"{namespace}/{graph_name}/{version}"]
+            if fn_name in graph.routers:
+                router_output = graph.invoke_router(fn_name, input)
+            else:
+                fn_output = graph.invoke_fn_ser(fn_name, input, init_value)
+
+                is_reducer = (
+                    graph.get_function(fn_name).indexify_function.accumulate is not None
+                )
+        except Exception as e:
+            print(traceback.format_exc())
+            has_failed = True
+            exception_msg = str(e)
+
+    # WARNING - IF THIS FAILS, WE WILL NOT BE ABLE TO RECOVER
+    # ANY LOGS
+    if has_failed:
+        return FunctionOutput(
+            fn_outputs=None,
+            router_output=None,
+            exception=exception_msg,
+            stdout=stdout_capture.getvalue(),
+            stderr=stderr_capture.getvalue(),
+            reducer=is_reducer,
+            success=False,
+        )
+    return FunctionOutput(
+        fn_outputs=fn_output,
+        router_output=router_output,
+        reducer=is_reducer,
+        success=True,
+        stdout=stdout_capture.getvalue(),
+        stderr=stderr_capture.getvalue(),
+    )

indexify/executor/indexify_executor.py

@@ -0,0 +1,32 @@
+import asyncio
+from typing import List, Optional
+
+import nanoid
+
+from .agent import ExtractorAgent
+from .function_worker import FunctionWorker
+
+
+def join(
+    workers: int,
+    server_addr: str = "localhost:8900",
+    config_path: Optional[str] = None,
+):
+    print(f"receiving tasks from server addr: {server_addr}")
+    id = nanoid.generate()
+    print(f"executor id: {id}")
+
+    function_worker = FunctionWorker(workers=workers)
+
+    agent = ExtractorAgent(
+        id,
+        num_workers=workers,
+        function_worker=function_worker,
+        server_addr=server_addr,
+        config_path=config_path,
+    )
+
+    try:
+        asyncio.get_event_loop().run_until_complete(agent.run())
+    except asyncio.CancelledError as ex:
+        print("exiting gracefully", ex)

indexify/executor/task_reporter.py

@@ -0,0 +1,110 @@
+import io
+from typing import List, Optional
+
+import httpx
+import nanoid
+from rich import print
+
+from indexify.executor.api_objects import RouterOutput as ApiRouterOutput
+from indexify.executor.api_objects import Task, TaskResult
+from indexify.executor.task_store import CompletedTask
+from indexify.functions_sdk.data_objects import IndexifyData, RouterOutput
+from indexify.functions_sdk.object_serializer import MsgPackSerializer
+
+
+# https://github.com/psf/requests/issues/1081#issuecomment-428504128
+class ForceMultipartDict(dict):
+    def __bool__(self):
+        return True
+
+
+FORCE_MULTIPART = ForceMultipartDict()
+
+
+class TaskReporter:
+    def __init__(self, base_url: str, executor_id: str):
+        self._base_url = base_url
+        self._executor_id = executor_id
+
+    def report_task_outcome(self, completed_task: CompletedTask):
+        fn_outputs = []
+        print(
+            f"[bold]task-reporter[/bold] uploading output of size: {len(completed_task.outputs or [])}"
+        )
+        for output in completed_task.outputs or []:
+            output_bytes = MsgPackSerializer.serialize(output)
+            fn_outputs.append(
+                ("node_outputs", (nanoid.generate(), io.BytesIO(output_bytes)))
+            )
+
+        if completed_task.errors:
+            print(
+                f"[bold]task-reporter[/bold] uploading error of size: {len(completed_task.errors)}"
+            )
+            fn_outputs.append(
+                (
+                    "exception_msg",
+                    (nanoid.generate(), io.BytesIO(completed_task.errors.encode())),
+                )
+            )
+
+        if completed_task.stdout:
+            print(
+                f"[bold]task-reporter[/bold] uploading stdout of size: {len(completed_task.stdout)}"
+            )
+            fn_outputs.append(
+                (
+                    "stdout",
+                    (nanoid.generate(), io.BytesIO(completed_task.stdout.encode())),
+                )
+            )
+
+        if completed_task.stderr:
+            print(
+                f"[bold]task-reporter[/bold] uploading stderr of size: {len(completed_task.stderr)}"
+            )
+            fn_outputs.append(
+                (
+                    "stderr",
+                    (nanoid.generate(), io.BytesIO(completed_task.stderr.encode())),
+                )
+            )
+
+        router_output = (
+            ApiRouterOutput(edges=completed_task.router_output.edges)
+            if completed_task.router_output
+            else None
+        )
+
+        task_result = TaskResult(
+            router_output=router_output,
+            outcome=completed_task.task_outcome,
+            namespace=completed_task.task.namespace,
+            compute_graph=completed_task.task.compute_graph,
+            compute_fn=completed_task.task.compute_fn,
+            invocation_id=completed_task.task.invocation_id,
+            executor_id=self._executor_id,
+            task_id=completed_task.task.id,
+            reducer=completed_task.reducer,
+        )
+        task_result_data = task_result.model_dump_json(exclude_none=True)
+
+        kwargs = {"data": {"task_result": task_result_data}}
+        if fn_outputs and len(fn_outputs) > 0:
+            kwargs["files"] = fn_outputs
+        else:
+            kwargs["files"] = FORCE_MULTIPART
+        try:
+            response = httpx.post(
+                url=f"{self._base_url}/internal/ingest_files",
+                **kwargs,
+            )
+        except Exception as e:
+            print(f"failed to report task outcome {e}")
+            raise e
+
+        try:
+            response.raise_for_status()
+        except Exception as e:
+            print(f"failed to report task outcome {response.text}")
+            raise e

indexify/executor/task_store.py

@@ -0,0 +1,113 @@
+import asyncio
+from typing import Dict, List, Literal, Optional
+
+from pydantic import BaseModel
+from rich import print
+
+from indexify.functions_sdk.data_objects import IndexifyData, RouterOutput
+
+from .api_objects import Task
+
+
+class CompletedTask(BaseModel):
+    task: Task
+    task_outcome: Literal["success", "failure"]
+    outputs: Optional[List[IndexifyData]] = None
+    router_output: Optional[RouterOutput] = None
+    errors: Optional[str] = None
+    stdout: Optional[str] = None
+    stderr: Optional[str] = None
+    reducer: bool = False
+
+
+class TaskStore:
+    def __init__(self) -> None:
+        self._tasks: Dict[str, Task] = {}
+        self._running_tasks: Dict[str, Task] = {}
+        self._finished: Dict[str, CompletedTask] = {}
+        self._retries: Dict[str, int] = {}
+        self._new_task_event = asyncio.Event()
+        self._finished_task_event = asyncio.Event()
+
+    def get_task(self, id) -> Task:
+        return self._tasks[id]
+
+    def add_tasks(self, tasks: List[Task]):
+        task: Task
+        for task in tasks:
+            if (
+                (task.id in self._tasks)
+                or (task.id in self._running_tasks)
+                or (task.id in self._finished)
+            ):
+                continue
+            print(
+                f"[bold] task store: [/bold] added task: {task.id} graph: {task.compute_graph} fn: {task.compute_fn} to queue"
+            )
+            self._tasks[task.id] = task
+            self._new_task_event.set()
+
+    async def get_runnable_tasks(self) -> Dict[str, Task]:
+        while True:
+            runnable_tasks = set(self._tasks) - set(self._running_tasks)
+            runnable_tasks = set(runnable_tasks) - set(self._finished)
+            if len(runnable_tasks) == 0:
+                await self._new_task_event.wait()
+                self._new_task_event.clear()
+            else:
+                break
+        out = {}
+        for task_id in runnable_tasks:
+            out[task_id] = self._tasks[task_id]
+            self._running_tasks[task_id] = self._tasks[task_id]
+        return out
+
+    def complete(self, outcome: CompletedTask):
+        self._retries.pop(outcome.task.id, None)
+        self._finished[outcome.task.id] = outcome
+        if outcome.task.id in self._running_tasks:
+            self._running_tasks.pop(outcome.task.id)
+        self._finished_task_event.set()
+
+    def retriable_failure(self, task_id: str):
+        self._running_tasks.pop(task_id)
+        if task_id not in self._retries:
+            self._retries[task_id] = 0
+        self._retries[task_id] += 1
+        if self._retries[task_id] > 3:
+            self._retries.pop(task_id)
+            self.complete(
+                outcome=CompletedTask(
+                    task_id=task_id, task_outcome="failed", outputs=[]
+                )
+            )
+        else:
+            self._new_task_event.set()
+
+    def mark_reported(self, task_id: str):
+        self._tasks.pop(task_id)
+        self._finished.pop(task_id)
+        print(f"[bold] task store: [/bold] removed task: {task_id} from queue")
+
+    def report_failed(self, task_id: str):
+        if self._finished[task_id].task_outcome != "Failed":
+            # An error occurred while reporting the task, mark it as failed
+            # and try reporting again.
+            self._finished[task_id].task_outcome = "Failed"
+        else:
+            # If a task is already marked as failed, remove it from the queue.
+            # The only possible error at this point is task not present at
+            # the coordinator.
+            self._tasks.pop(task_id)
+
+    def num_pending_tasks(self) -> int:
+        return len(self._tasks) + len(self._running_tasks)
+
+    async def task_outcomes(self) -> List[CompletedTask]:
+        while True:
+            if len(self._finished) == 0:
+                await self._finished_task_event.wait()
+                self._finished_task_event.clear()
+            else:
+                break
+        return self._finished.copy().values()

indexify/foo

@@ -0,0 +1,72 @@
+./functions_sdk/local_cache.py:        outputs = []
+./functions_sdk/local_cache.py:                outputs.append(f.read())
+./functions_sdk/local_cache.py:        return outputs
+./functions_sdk/graph.py:        outputs: List[Any] = fn_wrapper.run_fn(input, acc=acc)
+./functions_sdk/graph.py:            IndexifyData(payload=CborSerializer.serialize(output)) for output in outputs
+./functions_sdk/output_serializer.py:class OutputSerializer:
+./functions_sdk/output_serializer.py:        normalized_outputs = []
+./functions_sdk/output_serializer.py:            normalized_outputs.append(
+./functions_sdk/output_serializer.py:        return normalized_outputs
+./base_client.py:    def graph_outputs(
+./remote_client.py:class GraphOutputs(BaseModel):
+./remote_client.py:    outputs: List[GraphOutputMetadata]
+./remote_client.py:    def graph_outputs(
+./remote_client.py:            f"namespaces/{self.namespace}/compute_graphs/{graph}/invocations/{invocation_id}/outputs",
+./remote_client.py:        graph_outputs = GraphOutputs(**response.json())
+./remote_client.py:        outputs = []
+./remote_client.py:        for output in graph_outputs.outputs:
+./remote_client.py:                outputs.append(output)
+./remote_client.py:        return outputs
+./executor/task_reporter.py:        fn_outputs = []
+./executor/task_reporter.py:            f"[bold]task-reporter[/bold] uploading output of size: {len(completed_task.outputs)}"
+./executor/task_reporter.py:        for output in completed_task.outputs: 
+./executor/task_reporter.py:            fn_outputs.append(
+./executor/task_reporter.py:                ("node_outputs", (nanoid.generate(), io.BytesIO(output_bytes)))
+./executor/task_reporter.py:            fn_outputs.append(
+./executor/task_reporter.py:            fn_outputs.append(
+./executor/task_reporter.py:            fn_outputs.append(
+./executor/task_reporter.py:        if fn_outputs and len(fn_outputs) > 0:
+./executor/task_reporter.py:            kwargs["files"] = fn_outputs
+./executor/task_store.py:    outputs: List[IndexifyData]
+./executor/task_store.py:                    task_id=task_id, task_outcome="failed", outputs=[]
+./executor/downloader.py:            url = f"{self.base_url}/internal/fn_outputs/{task.input_key}"
+./executor/function_worker.py:    outputs: Union[List[IndexifyData], RouterOutput]
+./executor/function_worker.py:            indexify_data=result.outputs,
+./executor/function_worker.py:    return FunctionOutput(outputs=output, reducer=is_reducer)
+./executor/agent.py:                        f"Outputs: {len(task_outcome.outputs)}",
+./executor/agent.py:                            outputs=[],
+./executor/agent.py:                            outputs=[],
+./executor/agent.py:                            outputs=[],
+./executor/agent.py:                        outputs: FunctionWorkerOutput = await async_task
+./executor/agent.py:                            outputs if isinstance(outputs, RouterOutput) else None
+./executor/agent.py:                        if outputs.exception:
+./executor/agent.py:                            fn_outputs = []
+./executor/agent.py:                            fn_outputs = (
+./executor/agent.py:                                outputs.indexify_data if not isinstance(outputs, RouterOutput) else []
+./executor/agent.py:                            outputs=fn_outputs,
+./executor/agent.py:                            errors=outputs.exception,
+./executor/agent.py:                            stdout=outputs.stdout,
+./executor/agent.py:                            stderr=outputs.stderr,
+./executor/agent.py:                            reducer=outputs.reducer,
+./executor/agent.py:                            outputs=[],
+./local_client.py:# Holds the outputs of a
+./local_client.py:    outputs: Dict[str, List[IndexifyData]]
+./local_client.py:        outputs = defaultdict(list)
+./local_client.py:        self._results[input.id] = outputs
+./local_client.py:        self._run(g, input, outputs)
+./local_client.py:        outputs: Dict[str, List[bytes]],
+./local_client.py:                    f"ran {node_name}: num outputs: {len(cached_output_bytes)} (cache hit)"
+./local_client.py:                function_outputs: List[IndexifyData] = []
+./local_client.py:                    function_outputs.append(output)
+./local_client.py:                    outputs[node_name].append(output)
+./local_client.py:                function_outputs: List[IndexifyData] = g.invoke_fn_ser(
+./local_client.py:                print(f"ran {node_name}: num outputs: {len(function_outputs)}")
+./local_client.py:                    self._accumulators[node_name] = function_outputs[-1].model_copy()
+./local_client.py:                    outputs[node_name] = []
+./local_client.py:                outputs[node_name].extend(function_outputs)
+./local_client.py:                function_outputs_bytes: List[bytes] = [
+./local_client.py:                    for function_output in function_outputs
+./local_client.py:                    function_outputs_bytes,
+./local_client.py:                    for output in function_outputs:
+./local_client.py:                for output in function_outputs:
+./local_client.py:    def graph_outputs(

indexify/functions_sdk/data_objects.py

@@ -0,0 +1,37 @@
+from typing import Any, Dict, List, Optional, Union
+
+from pydantic import BaseModel, Json
+
+
+class FileInput(BaseModel):
+    url: str
+    mime_type: Optional[str] = None
+    metadata: Optional[Dict[str, Json]] = None
+    sha_256: Optional[str] = None
+
+
+class RouterOutput(BaseModel):
+    edges: List[str]
+
+
+class IndexifyData(BaseModel):
+    id: Optional[str] = None
+    payload: bytes
+    payload_encoding: str = "msgpack"
+
+
+class FunctionWorkerOutput(BaseModel):
+    fn_outputs: Optional[List[IndexifyData]]
+    router_output: Optional[RouterOutput]
+    exception: Optional[str]
+    stdout: Optional[str]
+    stderr: Optional[str]
+    reducer: bool = False
+    success: bool = True
+
+
+class File(BaseModel):
+    data: bytes
+    mime_type: Optional[str] = None
+    metadata: Optional[Dict[str, Any]] = None
+    sha_256: Optional[str] = None