indexify 0.2.42__tar.gz → 0.2.44__tar.gz

{indexify-0.2.42 → indexify-0.2.44}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: indexify
-Version: 0.2.42
+Version: 0.2.44
 Summary: Python Client for Indexify
 Home-page: https://github.com/tensorlakeai/indexify
 License: Apache 2.0
@@ -21,11 +21,11 @@ Requires-Dist: grpcio-tools (==1.68.1)
 Requires-Dist: httpx-sse (>=0.4.0,<0.5.0)
 Requires-Dist: httpx[http2] (==0.27.2)
 Requires-Dist: nanoid (>=2.0.0,<3.0.0)
-Requires-Dist: pydantic (==2.10.2)
+Requires-Dist: pydantic (>=2.10.2,<3.0.0)
 Requires-Dist: pyyaml (>=6,<7)
 Requires-Dist: rich (>=13.9.2,<14.0.0)
 Requires-Dist: structlog (>=24.4.0,<25.0.0)
-Requires-Dist: typer (>=0.13.0,<0.14.0)
+Requires-Dist: typer (>=0.12,<0.13)
 Project-URL: Repository, https://github.com/tensorlakeai/indexify
 Description-Content-Type: text/markdown
 

{indexify-0.2.42 → indexify-0.2.44}/indexify/cli.py

@@ -22,7 +22,7 @@ from rich.panel import Panel
 from rich.text import Text
 from rich.theme import Theme
 
-from indexify.executor.agent import ExtractorAgent
+from indexify.executor.executor import Executor
 from indexify.function_executor.function_executor_service import (
     FunctionExecutorService,
 )
@@ -221,7 +221,7 @@ def executor(
         shutil.rmtree(executor_cache)
     Path(executor_cache).mkdir(parents=True, exist_ok=True)
 
-    agent = ExtractorAgent(
+    executor = Executor(
         id,
         server_addr=server_addr,
         config_path=config_path,
@@ -232,7 +232,7 @@ def executor(
     )
 
     try:
-        asyncio.get_event_loop().run_until_complete(agent.run())
+        asyncio.get_event_loop().run_until_complete(executor.run())
     except asyncio.CancelledError:
         logger.info("graceful shutdown")
 

indexify-0.2.44/indexify/executor/executor.py

@@ -0,0 +1,155 @@
+import asyncio
+import signal
+from pathlib import Path
+from typing import Any, Optional
+
+import structlog
+
+from indexify.function_executor.proto.function_executor_pb2 import (
+    SerializedObject,
+)
+
+from .api_objects import Task
+from .downloader import DownloadedInputs, Downloader
+from .function_executor.process_function_executor_factory import (
+    ProcessFunctionExecutorFactory,
+)
+from .function_worker import (
+    FunctionWorker,
+    FunctionWorkerInput,
+    FunctionWorkerOutput,
+)
+from .task_fetcher import TaskFetcher
+from .task_reporter import TaskReporter
+
+
+class Executor:
+    def __init__(
+        self,
+        executor_id: str,
+        code_path: Path,
+        server_addr: str = "localhost:8900",
+        development_mode: bool = False,
+        config_path: Optional[str] = None,
+        name_alias: Optional[str] = None,
+        image_hash: Optional[str] = None,
+    ):
+        self._logger = structlog.get_logger(module=__name__)
+        self._should_run = True
+        self._config_path = config_path
+        protocol: str = "http"
+        if config_path:
+            self._logger.info("running the extractor with TLS enabled")
+            protocol = "https"
+
+        self._function_worker = FunctionWorker(
+            function_executor_factory=ProcessFunctionExecutorFactory(
+                indexify_server_address=server_addr,
+                development_mode=development_mode,
+                config_path=config_path,
+            )
+        )
+        self._server_addr = 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=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._task_reporter = TaskReporter(
+            base_url=self._base_url,
+            executor_id=executor_id,
+            config_path=self._config_path,
+        )
+
+    async def run(self):
+        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()
+        )
+
+        while self._should_run:
+            try:
+                async for task in self._task_fetcher.run():
+                    asyncio.create_task(self._run_task(task))
+            except Exception as e:
+                self._logger.error(
+                    "failed fetching tasks, retrying in 5 seconds", exc_info=e
+                )
+                await asyncio.sleep(5)
+
+    async def _run_task(self, task: Task) -> None:
+        """Runs the supplied task.
+
+        Doesn't raise any Exceptions. All errors are reported to the server."""
+        logger = self._task_logger(task)
+        output: Optional[FunctionWorkerOutput] = None
+
+        try:
+            graph: SerializedObject = await self._downloader.download_graph(task)
+            input: DownloadedInputs = await self._downloader.download_inputs(task)
+            output = await self._function_worker.run(
+                input=FunctionWorkerInput(
+                    task=task,
+                    graph=graph,
+                    function_input=input,
+                )
+            )
+            logger.info("task_execution_finished", success=output.success)
+        except Exception as e:
+            logger.error("failed running the task", exc_info=e)
+
+        await self._report_task_outcome(task=task, output=output, logger=logger)
+
+    async def _report_task_outcome(
+        self, task: Task, output: Optional[FunctionWorkerOutput], logger: Any
+    ) -> None:
+        """Reports the task with the given output to the server.
+
+        None output means that the task execution didn't finish due to an internal error.
+        Doesn't raise any exceptions."""
+        reporting_retries: int = 0
+
+        while True:
+            logger = logger.bind(retries=reporting_retries)
+            try:
+                await self._task_reporter.report(
+                    task=task, output=output, logger=logger
+                )
+                break
+            except Exception as e:
+                logger.error(
+                    "failed_to_report_task",
+                    exc_info=e,
+                )
+                reporting_retries += 1
+                await asyncio.sleep(5)
+
+    async def _shutdown(self, loop):
+        self._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):
+        loop.create_task(self._shutdown(loop))
+
+    def _task_logger(self, task: Task) -> Any:
+        return self._logger.bind(
+            namespace=task.namespace,
+            graph=task.compute_graph,
+            graph_version=task.graph_version,
+            invocation_id=task.invocation_id,
+            function_name=task.compute_fn,
+            task_id=task.id,
+        )

{indexify-0.2.42 → indexify-0.2.44}/indexify/executor/function_worker.py

@@ -31,13 +31,11 @@ class FunctionWorkerInput:
     def __init__(
         self,
         task: Task,
-        graph: Optional[SerializedObject] = None,
-        function_input: Optional[DownloadedInputs] = None,
+        graph: SerializedObject,
+        function_input: DownloadedInputs,
     ):
         self.task = task
-        # Must not be None when running the task.
         self.graph = graph
-        # Must not be None when running the task.
         self.function_input = function_input
 
 

{indexify-0.2.42 → indexify-0.2.44}/indexify/executor/task_fetcher.py

@@ -2,7 +2,6 @@ import json
 from importlib.metadata import version
 from typing import AsyncGenerator, Optional
 
-import httpx
 import structlog
 from httpx_sse import aconnect_sse
 
@@ -66,7 +65,7 @@ class TaskFetcher:
                 except Exception as e:
                     await event_source.response.aread()
                     raise Exception(
-                        "Failed to register at server. "
+                        "failed to register at server. "
                         f"Response code: {event_source.response.status_code}. "
                         f"Response text: '{event_source.response.text}'."
                     ) from e

indexify-0.2.44/indexify/executor/task_reporter.py

@@ -0,0 +1,215 @@
+import asyncio
+from typing import Any, List, Optional, Tuple
+
+import nanoid
+from httpx import Timeout
+
+from indexify.common_util import get_httpx_client
+from indexify.executor.api_objects import RouterOutput, Task, TaskResult
+from indexify.function_executor.proto.function_executor_pb2 import (
+    FunctionOutput,
+)
+
+from .function_worker import FunctionWorkerOutput
+
+
+# https://github.com/psf/requests/issues/1081#issuecomment-428504128
+class ForceMultipartDict(dict):
+    def __bool__(self):
+        return True
+
+
+FORCE_MULTIPART = ForceMultipartDict()
+UTF_8_CONTENT_TYPE = "application/octet-stream"
+
+
+class TaskOutputSummary:
+    def __init__(self):
+        self.output_count: int = 0
+        self.output_total_bytes: int = 0
+        self.router_output_count: int = 0
+        self.stdout_count: int = 0
+        self.stdout_total_bytes: int = 0
+        self.stderr_count: int = 0
+        self.stderr_total_bytes: int = 0
+        self.total_bytes: int = 0
+
+
+class TaskReporter:
+    def __init__(
+        self, base_url: str, executor_id: str, config_path: Optional[str] = None
+    ):
+        self._base_url = base_url
+        self._executor_id = executor_id
+        # Use thread-safe sync client due to issues with async client.
+        # Async client attempts to use connections it already closed.
+        # See e.g. https://github.com/encode/httpx/issues/2337.
+        # Creating a new async client for each request fixes this but it
+        # results in not reusing established TCP connections to server.
+        self._client = get_httpx_client(config_path, make_async=False)
+
+    async def report(
+        self, task: Task, output: Optional[FunctionWorkerOutput], logger: Any
+    ):
+        """Reports result of the supplied task.
+
+        If FunctionWorkerOutput is None this means that the task didn't finish and failed with internal error.
+        """
+        logger = logger.bind(module=__name__)
+        task_result, output_files, output_summary = self._process_task_output(
+            task, output
+        )
+        task_result_data = task_result.model_dump_json(exclude_none=True)
+
+        logger.info(
+            "reporting task outcome",
+            total_bytes=output_summary.total_bytes,
+            total_files=output_summary.output_count
+            + output_summary.stdout_count
+            + output_summary.stderr_count,
+            output_files=output_summary.output_count,
+            output_bytes=output_summary.total_bytes,
+            router_output_count=output_summary.router_output_count,
+            stdout_bytes=output_summary.stdout_total_bytes,
+            stderr_bytes=output_summary.stderr_total_bytes,
+        )
+
+        kwargs = {
+            "data": {"task_result": task_result_data},
+            # Use httpx default timeout of 5s for all timeout types.
+            # For read timeouts, use 5 minutes to allow for large file uploads.
+            "timeout": Timeout(
+                5.0,
+                read=5.0 * 60,
+            ),
+            "files": output_files if len(output_files) > 0 else FORCE_MULTIPART,
+        }
+        # Run in a separate thread to not block the main event loop.
+        response = await asyncio.to_thread(
+            self._client.post, url=f"{self._base_url}/internal/ingest_files", **kwargs
+        )
+
+        try:
+            response.raise_for_status()
+        except Exception as e:
+            # Caller catches and logs the exception.
+            raise Exception(
+                "failed to report task outcome. "
+                f"Response code: {response.status_code}. "
+                f"Response text: '{response.text}'."
+            ) from e
+
+    def _process_task_output(
+        self, task: Task, output: Optional[FunctionWorkerOutput]
+    ) -> Tuple[TaskResult, List[Any], TaskOutputSummary]:
+        task_result = TaskResult(
+            outcome="failure",
+            namespace=task.namespace,
+            compute_graph=task.compute_graph,
+            compute_fn=task.compute_fn,
+            invocation_id=task.invocation_id,
+            executor_id=self._executor_id,
+            task_id=task.id,
+        )
+        output_files: List[Any] = []
+        summary: TaskOutputSummary = TaskOutputSummary()
+        if output is None:
+            return task_result, output_files, summary
+
+        task_result.outcome = "success" if output.success else "failure"
+        task_result.reducer = output.reducer
+
+        _process_function_output(
+            function_output=output.function_output,
+            output_files=output_files,
+            summary=summary,
+        )
+        _process_router_output(
+            router_output=output.router_output, task_result=task_result, summary=summary
+        )
+        _process_stdout(
+            stdout=output.stdout, output_files=output_files, summary=summary
+        )
+        _process_stderr(
+            stderr=output.stderr, output_files=output_files, summary=summary
+        )
+
+        summary.total_bytes = (
+            summary.output_total_bytes
+            + summary.stdout_total_bytes
+            + summary.stderr_total_bytes
+        )
+
+        return task_result, output_files, summary
+
+
+def _process_function_output(
+    function_output: Optional[FunctionOutput],
+    output_files: List[Any],
+    summary: TaskOutputSummary,
+) -> None:
+    if function_output is None:
+        return
+
+    for output in function_output.outputs or []:
+        payload = output.bytes if output.HasField("bytes") else output.string
+        output_files.append(
+            (
+                "node_outputs",
+                (nanoid.generate(), payload, output.content_type),
+            )
+        )
+        summary.output_count += 1
+        summary.output_total_bytes += len(payload)
+
+
+def _process_router_output(
+    router_output: Optional[RouterOutput],
+    task_result: TaskResult,
+    summary: TaskOutputSummary,
+) -> None:
+    if router_output is None:
+        return
+
+    task_result.router_output = RouterOutput(edges=router_output.edges)
+    summary.router_output_count += 1
+
+
+def _process_stdout(
+    stdout: Optional[str], output_files: List[Any], summary: TaskOutputSummary
+) -> None:
+    if stdout is None:
+        return
+
+    output_files.append(
+        (
+            "stdout",
+            (
+                nanoid.generate(),
+                stdout.encode(),
+                UTF_8_CONTENT_TYPE,
+            ),
+        )
+    )
+    summary.stdout_count += 1
+    summary.stdout_total_bytes += len(stdout)
+
+
+def _process_stderr(
+    stderr: Optional[str], output_files: List[Any], summary: TaskOutputSummary
+) -> None:
+    if stderr is None:
+        return
+
+    output_files.append(
+        (
+            "stderr",
+            (
+                nanoid.generate(),
+                stderr.encode(),
+                UTF_8_CONTENT_TYPE,
+            ),
+        )
+    )
+    summary.stderr_count += 1
+    summary.stderr_total_bytes += len(stderr)

{indexify-0.2.42 → indexify-0.2.44}/indexify/functions_sdk/graph.py

@@ -37,7 +37,6 @@ from .indexify_functions import (
     IndexifyRouter,
     RouterCallResult,
 )
-from .local_cache import CacheAwareFunctionWrapper
 from .object_serializer import get_serializer
 
 RouterFn = Annotated[
@@ -86,7 +85,6 @@ class Graph:
 
         # Storage for local execution
         self._results: Dict[str, Dict[str, List[IndexifyData]]] = {}
-        self._cache = CacheAwareFunctionWrapper("./indexify_local_runner_cache")
         self._accumulator_values: Dict[str, IndexifyData] = {}
         self._local_graph_ctx: Optional[GraphInvocationContext] = None
 

{indexify-0.2.42 → indexify-0.2.44}/pyproject.toml

@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "indexify"
-version = "0.2.42"
+version = "0.2.44"
 description = "Python Client for Indexify"
 authors = ["Tensorlake Inc. <support@tensorlake.ai>"]
 license = "Apache 2.0"
@@ -15,12 +15,12 @@ indexify-cli = "indexify.cli:app"
 python = "^3.9"
 httpx = { version = "0.27.2", extras = ["http2"] }
 pyyaml = "^6"
-pydantic = "2.10.2"
+pydantic = "^2.10.2"
 cloudpickle = "^3.1.0"
 rich = "^13.9.2"
 nanoid = "^2.0.0"
 docker = "^7.1.0"
-typer = "^0.13.0"
+typer = "^0.12"
 httpx-sse = "^0.4.0"
 structlog = "^24.4.0"
 grpcio = "1.68.1"

indexify-0.2.42/indexify/executor/agent.py

@@ -1,262 +0,0 @@
-import asyncio
-from pathlib import Path
-from typing import Dict, List, Optional
-
-import structlog
-
-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,
-)
-from .task_fetcher import TaskFetcher
-from .task_reporter import TaskReporter
-from .task_store import CompletedTask, TaskStore
-
-logger = structlog.get_logger(module=__name__)
-
-
-class ExtractorAgent:
-    def __init__(
-        self,
-        executor_id: str,
-        code_path: Path,
-        server_addr: str = "localhost:8900",
-        development_mode: bool = False,
-        config_path: Optional[str] = None,
-        name_alias: Optional[str] = None,
-        image_hash: Optional[str] = None,
-    ):
-        self._config_path = config_path
-        protocol: str = "http"
-        if config_path:
-            logger.info("running the extractor with TLS enabled")
-            protocol = "https"
-
-        self._task_store: TaskStore = TaskStore()
-        self._function_worker = FunctionWorker(
-            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"{protocol}://{self._server_addr}"
-        self._code_path = code_path
-        self._downloader = Downloader(
-            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._task_reporter = TaskReporter(
-            base_url=self._base_url,
-            executor_id=executor_id,
-            config_path=self._config_path,
-        )
-
-    async def task_completion_reporter(self):
-        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:
-                logger.info(
-                    "reporting_task_outcome",
-                    task_id=task_outcome.task.id,
-                    fn_name=task_outcome.task.compute_fn,
-                    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,
-                )
-
-                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
-                    logger.error(
-                        "failed_to_report_task",
-                        task_id=task_outcome.task.id,
-                        exc_info=e,
-                        retries=task_outcome.reporting_retries,
-                    )
-                    task_outcome.reporting_retries += 1
-                    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] = [
-            asyncio.create_task(
-                self._task_store.get_runnable_tasks(), name="get_runnable_tasks"
-            )
-        ]
-
-        while True:
-            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():
-                        logger.error(
-                            "task_launcher_error, failed to get runnable tasks",
-                            exc_info=async_task.exception(),
-                        )
-                        continue
-                    result: Dict[str, Task] = await async_task
-                    task: Task
-                    for _, task in result.items():
-                        async_tasks.append(
-                            DownloadGraphTask(
-                                function_worker_input=FunctionWorkerInput(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():
-                        logger.error(
-                            "task_launcher_error, failed to download graph",
-                            exc_info=async_task.exception(),
-                        )
-                        completed_task = CompletedTask(
-                            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(
-                        DownloadInputsTask(
-                            function_worker_input=function_worker_input,
-                            downloader=self._downloader,
-                        )
-                    )
-                elif async_task.get_name() == "download_inputs":
-                    if async_task.exception():
-                        logger.error(
-                            "task_launcher_error, failed to download inputs",
-                            exc_info=async_task.exception(),
-                        )
-                        completed_task = CompletedTask(
-                            task=async_task.function_worker_input.task,
-                            task_outcome="failure",
-                        )
-                        self._task_store.complete(outcome=completed_task)
-                        continue
-                    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_task":
-                    if async_task.exception():
-                        completed_task = CompletedTask(
-                            task=async_task.function_worker_input.task,
-                            task_outcome="failure",
-                            stderr=str(async_task.exception()),
-                        )
-                        self._task_store.complete(outcome=completed_task)
-                        continue
-                    async_task: RunTask
-                    try:
-                        outputs: FunctionWorkerOutput = await async_task
-                        if not outputs.success:
-                            task_outcome = "failure"
-                        else:
-                            task_outcome = "success"
-
-                        completed_task = CompletedTask(
-                            task=async_task.function_worker_input.task,
-                            task_outcome=task_outcome,
-                            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 Exception as e:
-                        logger.error(
-                            "failed to execute task",
-                            task_id=async_task.function_worker_input.task.id,
-                            exc_info=e,
-                        )
-                        completed_task = CompletedTask(
-                            task=async_task.function_worker_input.task,
-                            task_outcome="failure",
-                        )
-                        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())
-        await self._main_loop()
-
-    async def _shutdown(self, loop):
-        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):
-        loop.create_task(self._shutdown(loop))

indexify-0.2.42/indexify/executor/executor_tasks.py

@@ -1,58 +0,0 @@
-import asyncio
-
-from pydantic import BaseModel
-
-from .api_objects import Task
-from .downloader import Downloader
-from .function_worker import FunctionWorker, FunctionWorkerInput
-
-
-class DownloadGraphTask(asyncio.Task):
-    def __init__(
-        self,
-        *,
-        function_worker_input: FunctionWorkerInput,
-        downloader: Downloader,
-        **kwargs,
-    ):
-        kwargs["name"] = "download_graph"
-        kwargs["loop"] = asyncio.get_event_loop()
-        super().__init__(
-            downloader.download_graph(function_worker_input.task),
-            **kwargs,
-        )
-        self.function_worker_input = function_worker_input
-
-
-class DownloadInputsTask(asyncio.Task):
-    def __init__(
-        self,
-        *,
-        function_worker_input: FunctionWorkerInput,
-        downloader: Downloader,
-        **kwargs,
-    ):
-        kwargs["name"] = "download_inputs"
-        kwargs["loop"] = asyncio.get_event_loop()
-        super().__init__(
-            downloader.download_inputs(function_worker_input.task),
-            **kwargs,
-        )
-        self.function_worker_input = function_worker_input
-
-
-class RunTask(asyncio.Task):
-    def __init__(
-        self,
-        *,
-        function_worker: FunctionWorker,
-        function_worker_input: FunctionWorkerInput,
-        **kwargs,
-    ):
-        kwargs["name"] = "run_task"
-        kwargs["loop"] = asyncio.get_event_loop()
-        super().__init__(
-            function_worker.run(function_worker_input),
-            **kwargs,
-        )
-        self.function_worker_input = function_worker_input

indexify-0.2.42/indexify/executor/task_reporter.py

@@ -1,155 +0,0 @@
-from typing import Optional
-
-import nanoid
-import structlog
-from httpx import Timeout
-from pydantic import BaseModel
-
-from indexify.common_util import get_httpx_client
-from indexify.executor.api_objects import RouterOutput as ApiRouterOutput
-from indexify.executor.api_objects import TaskResult
-from indexify.executor.task_store import CompletedTask
-
-logger = structlog.get_logger(__name__)
-
-
-# https://github.com/psf/requests/issues/1081#issuecomment-428504128
-class ForceMultipartDict(dict):
-    def __bool__(self):
-        return True
-
-
-FORCE_MULTIPART = ForceMultipartDict()
-UTF_8_CONTENT_TYPE = "application/octet-stream"
-
-
-class ReportingData(BaseModel):
-    output_count: int = 0
-    output_total_bytes: int = 0
-    stdout_count: int = 0
-    stdout_total_bytes: int = 0
-    stderr_count: int = 0
-    stderr_total_bytes: int = 0
-
-
-class TaskReporter:
-    def __init__(
-        self, base_url: str, executor_id: str, config_path: Optional[str] = None
-    ):
-        self._base_url = base_url
-        self._executor_id = executor_id
-        self._client = get_httpx_client(config_path)
-
-    def report_task_outcome(self, completed_task: CompletedTask):
-        report = ReportingData()
-        fn_outputs = []
-
-        if completed_task.function_output:
-            for output in completed_task.function_output.outputs or []:
-                payload = output.bytes if output.HasField("bytes") else output.string
-                fn_outputs.append(
-                    (
-                        "node_outputs",
-                        (nanoid.generate(), payload, output.content_type),
-                    )
-                )
-                report.output_count += 1
-                report.output_total_bytes += len(payload)
-
-        if completed_task.stdout:
-            fn_outputs.append(
-                (
-                    "stdout",
-                    (
-                        nanoid.generate(),
-                        completed_task.stdout.encode(),
-                        UTF_8_CONTENT_TYPE,
-                    ),
-                )
-            )
-            report.stdout_count += 1
-            report.stdout_total_bytes += len(completed_task.stdout)
-
-        if completed_task.stderr:
-            fn_outputs.append(
-                (
-                    "stderr",
-                    (
-                        nanoid.generate(),
-                        completed_task.stderr.encode(),
-                        UTF_8_CONTENT_TYPE,
-                    ),
-                )
-            )
-            report.stderr_count += 1
-            report.stderr_total_bytes += len(completed_task.stderr)
-
-        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)
-
-        total_bytes = (
-            report.output_total_bytes
-            + report.stdout_total_bytes
-            + report.stderr_total_bytes
-        )
-
-        logger.info(
-            "reporting task outcome",
-            task_id=completed_task.task.id,
-            retries=completed_task.reporting_retries,
-            total_bytes=total_bytes,
-            total_files=report.output_count + report.stdout_count + report.stderr_count,
-            output_files=report.output_count,
-            output_bytes=total_bytes,
-            stdout_bytes=report.stdout_total_bytes,
-            stderr_bytes=report.stderr_total_bytes,
-        )
-        #
-        kwargs = {
-            "data": {"task_result": task_result_data},
-            # Use httpx default timeout of 5s for all timeout types.
-            # For read timeouts, use 5 minutes to allow for large file uploads.
-            "timeout": Timeout(
-                5.0,
-                read=5.0 * 60,
-            ),
-        }
-        if fn_outputs and len(fn_outputs) > 0:
-            kwargs["files"] = fn_outputs
-        else:
-            kwargs["files"] = FORCE_MULTIPART
-
-        response = self._client.post(
-            url=f"{self._base_url}/internal/ingest_files",
-            **kwargs,
-        )
-
-        try:
-            response.raise_for_status()
-        except Exception as e:
-            # Caller catches and logs the exception.
-            # Log response details here for easier debugging.
-            logger.error(
-                "failed to report task outcome",
-                task_id=completed_task.task.id,
-                retries=completed_task.reporting_retries,
-                status_code=response.status_code,
-                response_text=response.text,
-            )
-            raise e

indexify-0.2.42/indexify/executor/task_store.py

@@ -1,132 +0,0 @@
-import asyncio
-from typing import Dict, List, Literal, Optional
-
-import structlog
-
-from indexify.function_executor.proto.function_executor_pb2 import (
-    FunctionOutput,
-    RouterOutput,
-)
-
-from .api_objects import Task
-
-logger = structlog.get_logger(module=__name__)
-
-
-class CompletedTask:
-    def __init__(
-        self,
-        task: Task,
-        task_outcome: Literal["success", "failure"],
-        function_output: Optional[FunctionOutput] = None,
-        router_output: Optional[RouterOutput] = None,
-        stdout: Optional[str] = None,
-        stderr: Optional[str] = None,
-        reducer: bool = False,
-        reporting_retries: int = 0,
-    ):
-        self.task = task
-        self.task_outcome = task_outcome
-        self.function_output = function_output
-        self.router_output = router_output
-        self.stdout = stdout
-        self.stderr = stderr
-        self.reducer = reducer
-        self.reporting_retries = reporting_retries
-
-
-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
-            logger.info(
-                "added task",
-                task_id=task.id,
-                namespace=task.namespace,
-                graph=task.compute_graph,
-                fn=task.compute_fn,
-            )
-            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)
-        logger.info("removed task", task_id=task_id)
-
-    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-0.2.42/indexify/functions_sdk/local_cache.py

@@ -1,46 +0,0 @@
-import os
-from hashlib import sha256
-from typing import List, Optional
-
-
-class CacheAwareFunctionWrapper:
-    def __init__(self, cache_dir: str):
-        self._cache_dir = cache_dir
-        if not os.path.exists(cache_dir):
-            os.makedirs(cache_dir)
-
-    def _get_key(self, input: bytes) -> str:
-        h = sha256()
-        h.update(input)
-        return h.hexdigest()
-
-    def get(self, graph: str, node_name: str, input: bytes) -> Optional[List[bytes]]:
-        key = self._get_key(input)
-        dir_path = os.path.join(self._cache_dir, graph, node_name, key)
-        if not os.path.exists(dir_path):
-            return None
-
-        files = os.listdir(dir_path)
-        outputs = []
-        for file in files:
-            with open(os.path.join(dir_path, file), "rb") as f:
-                return f.read()
-
-        return outputs
-
-    def set(
-        self,
-        graph: str,
-        node_name: str,
-        input: bytes,
-        output: List[bytes],
-    ):
-        key = self._get_key(input)
-        dir_path = os.path.join(self._cache_dir, graph, node_name, key)
-        if not os.path.exists(dir_path):
-            os.makedirs(dir_path)
-
-        for i, output_item in enumerate(output):
-            file_path = os.path.join(dir_path, f"{i}.cbor")
-            with open(file_path, "wb") as f:
-                f.write(output_item)