python-sdk(feat): Support async function #1017

.gitignore

@@ -27,6 +27,7 @@ charts/*.tgz
 indexify_storage
 indexify_local_runner_cache
 server/indexify_storage
+server/indexify_server_state
 local_cache
 .dev-tls
 src/state/store/snapshots/*

python-sdk/indexify/cli.py

@@ -16,7 +16,6 @@
 from rich.theme import Theme
 
 from indexify.executor.agent import ExtractorAgent
-from indexify.executor.function_worker import FunctionWorker
 from indexify.functions_sdk.image import (
     DEFAULT_IMAGE_3_10,
     DEFAULT_IMAGE_3_11,
@@ -31,7 +30,6 @@
         "highlight": "magenta",
     }
 )
-
 console = Console(theme=custom_theme)
 
 app = typer.Typer(pretty_exceptions_enable=False, no_args_is_help=True)
@@ -250,7 +248,6 @@ def _build_image(image: Image, python_sdk_path: Optional[str] = None):
 
     docker_file += "\n".join(run_strs)
     print(os.getcwd())
-    import docker
     import docker.api.build
 
     docker.api.build.process_dockerfile = lambda dockerfile, path: (

python-sdk/indexify/executor/api_objects.py

@@ -29,10 +29,6 @@ class RouterOutput(BaseModel):
     edges: List[str]
 
 
-class FnOutput(BaseModel):
-    payload: Json
-
-
 class TaskResult(BaseModel):
     router_output: Optional[RouterOutput] = None
     outcome: str

python-sdk/indexify/executor/downloader.py

@@ -1,17 +1,18 @@
+import asyncio
 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 ..common_util import get_httpx_client
+from .. import IndexifyClient
 from ..functions_sdk.object_serializer import JsonSerializer, get_serializer
 from .api_objects import Task
+from .executor_tasks import DownloadTask, TaskEnum
 
 custom_theme = Theme(
     {
@@ -31,11 +32,29 @@ class DownloadedInputs(BaseModel):
 
 class Downloader:
     def __init__(
-        self, code_path: str, base_url: str, config_path: Optional[str] = None
+        self,
+        code_path: str,
+        base_url: str,
+        indexify_client: IndexifyClient,
     ):
         self.code_path = code_path
         self.base_url = base_url
-        self._client = get_httpx_client(config_path)
+        self._indexify_client = indexify_client
+        self._event_loop = asyncio.get_event_loop()
+
+    def download(self, task, name):
+        if name == TaskEnum.DOWNLOAD_GRAPH_TASK:
+            coroutine = self.download_graph(
+                task.namespace, task.compute_graph, task.graph_version
+            )
+        elif name == TaskEnum.DOWNLOAD_INPUT_TASK:
+            coroutine = self.download_input(task)
+        else:
+            raise ValueError(f"Unsupported task name: {name}")
+
+        return DownloadTask(
+            task=task, coroutine=coroutine, name=name, loop=self._event_loop
+        )
 
     async def download_graph(self, namespace: str, name: str, version: int) -> str:
         path = os.path.join(self.code_path, namespace, f"{name}.{version}")
@@ -50,21 +69,7 @@ async def download_graph(self, namespace: str, name: str, version: int) -> str:
             )
         )
 
-        response = self._client.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
-
+        response = self._indexify_client.download_graph(namespace, name)
         os.makedirs(os.path.dirname(path), exist_ok=True)
         with open(path, "wb") as f:
             f.write(response.content)
@@ -83,25 +88,29 @@ async def download_input(self, task: Task) -> DownloadedInputs:
 
         console.print(
             Panel(
-                f"downloading input\nURL: {url} \n reducer input URL: {reducer_url}",
+                f"downloading input\nFunction: {task.compute_fn} \n reducer id: {task.reducer_output_id}",
                 title="downloader",
                 border_style="cyan",
             )
         )
 
-        response = self._client.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",
-                )
+        input_id = task.input_key.split("|")[-1]
+        if task.invocation_id == input_id:
+            response = self._indexify_client.download_fn_input(
+                task.namespace, task.compute_graph, task.invocation_id
+            )
+        else:
+            response = self._indexify_client.download_fn_output(task.input_key)
+
+        init_value = None
+        if task.reducer_output_id:
+            init_value = self._indexify_client.download_reducer_input(
+                task.namespace,
+                task.compute_graph,
+                task.invocation_id,
+                task.compute_fn,
+                task.reducer_output_id,
             )
-            raise
 
         encoder = (
             "json"
@@ -119,19 +128,7 @@ async def download_input(self, task: Task) -> DownloadedInputs:
 
         deserialized_content = serializer.deserialize(response.content)
 
-        if reducer_url:
-            init_value = self._client.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
+        if init_value:
             init_value = serializer.deserialize(init_value.content)
             return DownloadedInputs(
                 input=IndexifyData(

python-sdk/indexify/executor/executor_tasks.py

@@ -1,73 +1,59 @@
 import asyncio
-from typing import Optional
-
-from indexify.functions_sdk.data_objects import IndexifyData
+from enum import Enum, unique
+from typing import Coroutine
 
 from .api_objects import Task
-from .downloader import Downloader
-from .function_worker import FunctionWorker
 
 
-class DownloadGraphTask(asyncio.Task):
+class RunFunctionTask(asyncio.Task):
     def __init__(
         self,
         *,
         task: Task,
-        downloader: Downloader,
+        coroutine: Coroutine,
+        loop: asyncio.AbstractEventLoop,
         **kwargs,
     ):
-        kwargs["name"] = "download_graph"
-        kwargs["loop"] = asyncio.get_event_loop()
+        kwargs["name"] = TaskEnum.RUN_FUNCTION_TASK.value
+        kwargs["loop"] = loop
         super().__init__(
-            downloader.download_graph(
-                task.namespace, task.compute_graph, task.graph_version
-            ),
+            coroutine,
             **kwargs,
         )
         self.task = task
 
 
-class DownloadInputTask(asyncio.Task):
+class DownloadTask(asyncio.Task):
     def __init__(
         self,
         *,
         task: Task,
-        downloader: Downloader,
+        coroutine: Coroutine,
+        name: str,
+        loop: asyncio.AbstractEventLoop,
         **kwargs,
     ):
-        kwargs["name"] = "download_input"
-        kwargs["loop"] = asyncio.get_event_loop()
+        if not isinstance(name, TaskEnum):
+            raise ValueError(f"name '{name}' must be TaskEnum")
+        kwargs["name"] = name.value
+        kwargs["loop"] = loop
         super().__init__(
-            downloader.download_input(task),
+            coroutine,
             **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,
-                invocation_id=task.invocation_id,
-            ),
-            **kwargs,
-        )
-        self.task = task
+@unique
+class TaskEnum(Enum):
+    GET_RUNNABLE_TASK = "get_runnable_tasks"
+    RUN_FUNCTION_TASK = "run_function"
+    DOWNLOAD_GRAPH_TASK = "download_graph"
+    DOWNLOAD_INPUT_TASK = "download_input"
+
+    @classmethod
+    def from_value(cls, value):
+        for task_name in cls:
+            if task_name.value == value:
+                return task_name
+        raise ValueError(f"No task found with value {value}")

python-sdk/indexify/executor/function_worker/__init__.py

@@ -0,0 +1,3 @@
+from .function_worker import FunctionWorker
+
+__all__ = ["FunctionWorker"]