indexify 0.2.22__py3-none-any.whl → 0.2.24__py3-none-any.whl

indexify/__init__.py

@@ -2,6 +2,8 @@ from . import data_loaders
 from .functions_sdk.graph import Graph
 from .functions_sdk.image import Image
 from .functions_sdk.indexify_functions import (
+    IndexifyFunction,
+    get_ctx,
     indexify_function,
     indexify_router,
 )
@@ -19,6 +21,8 @@ __all__ = [
     "RemotePipeline",
     "Image",
     "indexify_function",
+    "get_ctx",
+    "IndexifyFunction",
     "indexify_router",
     "DEFAULT_SERVICE_URL",
     "IndexifyClient",

indexify/cli.py

@@ -170,6 +170,9 @@ def executor(
     name_alias: Optional[str] = typer.Option(
         None, help="Name alias for the executor if it's spun up with the base image"
     ),
+    image_version: Optional[int] = typer.Option(
+        "1", help="Requested Image Version for this executor"
+    ),
 ):
     id = nanoid.generate()
     console.print(
@@ -179,13 +182,13 @@ def executor(
             f"Server address: {server_addr}\n"
             f"Executor ID: {id}\n"
             f"Executor cache: {executor_cache}\n"
-            f"Name Alias: {name_alias}",
+            f"Name Alias: {name_alias}"
+            f"Image Version: {image_version}\n",
             title="Agent Configuration",
             border_style="info",
         )
     )
 
-    function_worker = FunctionWorker(workers=workers)
     from pathlib import Path
 
     executor_cache = Path(executor_cache).expanduser().absolute()
@@ -196,11 +199,11 @@ def executor(
     agent = ExtractorAgent(
         id,
         num_workers=workers,
-        function_worker=function_worker,
         server_addr=server_addr,
         config_path=config_path,
         code_path=executor_cache,
         name_alias=name_alias,
+        image_version=image_version,
     )
 
     try:

indexify/error.py

@@ -5,4 +5,4 @@ class ApiException(Exception):
 
 class GraphStillProcessing(Exception):
     def __init__(self) -> None:
-        super().__init__("graph is still processing")
+        super().__init__("graph is still processing")

indexify/executor/agent.py

@@ -1,14 +1,12 @@
 import asyncio
 import json
 import ssl
+import traceback
 from concurrent.futures.process import BrokenProcessPool
 from importlib.metadata import version
-import traceback
 from typing import Dict, List, Optional
 
 import httpx
-from indexify.functions_sdk.graph_definition import ComputeGraphMetadata
-from indexify.http_client import IndexifyClient
 import yaml
 from httpx_sse import aconnect_sse
 from pydantic import BaseModel
@@ -20,10 +18,12 @@ from rich.theme import Theme
 from indexify.functions_sdk.data_objects import (
     FunctionWorkerOutput,
     IndexifyData,
-    RouterOutput,
 )
-from . import image_dependency_installer
+from indexify.functions_sdk.graph_definition import ComputeGraphMetadata
+from indexify.http_client import IndexifyClient
 
+from ..functions_sdk.image import ImageInformation
+from . import image_dependency_installer
 from .api_objects import ExecutorMetadata, Task
 from .downloader import DownloadedInputs, Downloader
 from .executor_tasks import DownloadGraphTask, DownloadInputTask, ExtractTask
@@ -31,7 +31,6 @@ from .function_worker import FunctionWorker
 from .runtime_probes import ProbeInfo, RuntimeProbes
 from .task_reporter import TaskReporter
 from .task_store import CompletedTask, TaskStore
-from ..functions_sdk.image import ImageInformation
 
 custom_theme = Theme(
     {
@@ -60,12 +59,13 @@ class ExtractorAgent:
         executor_id: str,
         num_workers,
         code_path: str,
-        function_worker: FunctionWorker,
         server_addr: str = "localhost:8900",
         config_path: Optional[str] = None,
         name_alias: Optional[str] = None,
+        image_version: Optional[int] = None,
     ):
         self.name_alias = name_alias
+        self.image_version = image_version
 
         self._probe = RuntimeProbes()
 
@@ -111,7 +111,12 @@ class ExtractorAgent:
 
         self._task_store: TaskStore = TaskStore()
         self._executor_id = executor_id
-        self._function_worker = function_worker
+        self._function_worker = FunctionWorker(
+            workers=num_workers,
+            indexify_client=IndexifyClient(
+                service_url=f"{self._protocol}://{server_addr}"
+            ),
+        )
         self._has_registered = False
         self._server_addr = server_addr
         self._base_url = f"{self._protocol}://{self._server_addr}"
@@ -195,7 +200,7 @@ class ExtractorAgent:
                             task, self._protocol, self._server_addr
                         )
                         image_dependency_installer.executor_image_builder(
-                            image_info, self.name_alias
+                            image_info, self.name_alias, self.image_version
                         )
                         self._require_image_bootstrap = False
                     except Exception as e:
@@ -309,7 +314,7 @@ class ExtractorAgent:
                             task=async_task.task,
                             task_outcome="failure",
                             outputs=[],
-                            errors=str(async_task.exception()),
+                            stderr=str(async_task.exception()),
                         )
                         self._task_store.complete(outcome=completed_task)
                         continue
@@ -326,7 +331,6 @@ class ExtractorAgent:
                             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,
@@ -378,11 +382,18 @@ class ExtractorAgent:
                 else runtime_probe.image_name
             )
 
+            image_version: int = (
+                self.image_version
+                if self.image_version is not None
+                else runtime_probe.image_version
+            )
+
             data = ExecutorMetadata(
                 id=self._executor_id,
                 executor_version=executor_version,
                 addr="",
                 image_name=image_name,
+                image_version=image_version,
                 labels=runtime_probe.labels,
             ).model_dump()
 
@@ -406,6 +417,11 @@ class ExtractorAgent:
                         json=data,
                         headers={"Content-Type": "application/json"},
                     ) as event_source:
+                        if not event_source.response.is_success:
+                            resp = await event_source.response.aread().decode('utf-8')
+                            console.print(f"failed to register: {str(resp)}")
+                            await asyncio.sleep(5)
+                            continue
                         console.print(
                             Text("executor registered successfully", style="bold green")
                         )

indexify/executor/api_objects.py

@@ -21,6 +21,7 @@ class ExecutorMetadata(BaseModel):
     executor_version: str
     addr: str
     image_name: str
+    image_version: int
     labels: Dict[str, Any]
 
 

indexify/executor/executor_tasks.py

@@ -66,6 +66,7 @@ class ExtractTask(asyncio.Task):
                 init_value=init_value,
                 code_path=code_path,
                 version=task.graph_version,
+                invocation_id=task.invocation_id,
             ),
             **kwargs,
         )

indexify/executor/function_worker.py

@@ -6,6 +6,7 @@ import cloudpickle
 from pydantic import BaseModel
 from rich import print
 
+from indexify import IndexifyClient
 from indexify.functions_sdk.data_objects import (
     FunctionWorkerOutput,
     IndexifyData,
@@ -13,6 +14,7 @@ from indexify.functions_sdk.data_objects import (
 )
 from indexify.functions_sdk.indexify_functions import (
     FunctionCallResult,
+    GraphInvocationContext,
     IndexifyFunctionWrapper,
     RouterCallResult,
 )
@@ -38,13 +40,18 @@ class FunctionOutput(BaseModel):
     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
+    namespace: str,
+    graph_name: str,
+    fn_name: str,
+    code_path: str,
+    version: int,
+    invocation_id: str,
+    indexify_client: IndexifyClient,
 ):
     """Load an extractor to the memory: extractor_wrapper_map."""
     global function_wrapper_map
@@ -54,18 +61,28 @@ def _load_function(
     with open(code_path, "rb") as f:
         code = f.read()
         pickled_functions = cloudpickle.loads(code)
+    context = GraphInvocationContext(
+        invocation_id=invocation_id,
+        graph_name=graph_name,
+        graph_version=str(version),
+        indexify_client=indexify_client,
+    )
     function_wrapper = IndexifyFunctionWrapper(
-        cloudpickle.loads(pickled_functions[fn_name])
+        cloudpickle.loads(pickled_functions[fn_name]),
+        context,
     )
     function_wrapper_map[key] = function_wrapper
 
 
 class FunctionWorker:
-    def __init__(self, workers: int = 1) -> None:
+    def __init__(
+        self, workers: int = 1, indexify_client: IndexifyClient = None
+    ) -> None:
         self._executor: concurrent.futures.ProcessPoolExecutor = (
             concurrent.futures.ProcessPoolExecutor(max_workers=workers)
         )
         self._workers = workers
+        self._indexify_client = indexify_client
 
     async def async_submit(
         self,
@@ -76,15 +93,23 @@ class FunctionWorker:
         code_path: str,
         version: int,
         init_value: Optional[IndexifyData] = None,
+        invocation_id: Optional[str] = None,
     ) -> FunctionWorkerOutput:
         try:
             result = _run_function(
-                namespace, graph_name, fn_name, input, code_path, version, init_value
+                namespace,
+                graph_name,
+                fn_name,
+                input,
+                code_path,
+                version,
+                init_value,
+                invocation_id,
+                self._indexify_client,
             )
             # TODO - bring back running in a separate process
         except Exception as e:
             return FunctionWorkerOutput(
-                exception=str(e),
                 stdout=e.stdout,
                 stderr=e.stderr,
                 reducer=e.is_reducer,
@@ -94,7 +119,6 @@ class FunctionWorker:
         return FunctionWorkerOutput(
             fn_outputs=result.fn_outputs,
             router_output=result.router_output,
-            exception=result.exception,
             stdout=result.stdout,
             stderr=result.stderr,
             reducer=result.reducer,
@@ -113,6 +137,8 @@ def _run_function(
     code_path: str,
     version: int,
     init_value: Optional[IndexifyData] = None,
+    invocation_id: Optional[str] = None,
+    indexify_client: Optional[IndexifyClient] = None,
 ) -> FunctionOutput:
     import io
     from contextlib import redirect_stderr, redirect_stdout
@@ -123,7 +149,6 @@ def _run_function(
     router_output = None
     fn_output = None
     has_failed = False
-    exception_msg = None
     print(
         f"[bold] function_worker: [/bold] invoking function {fn_name} in graph {graph_name}"
     )
@@ -131,7 +156,15 @@ def _run_function(
         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)
+                _load_function(
+                    namespace,
+                    graph_name,
+                    fn_name,
+                    code_path,
+                    version,
+                    invocation_id,
+                    indexify_client,
+                )
 
             fn = function_wrapper_map[key]
             if (
@@ -143,7 +176,6 @@ def _run_function(
                 if router_call_result.traceback_msg is not None:
                     print(router_call_result.traceback_msg, file=sys.stderr)
                     has_failed = True
-                    exception_msg = router_call_result.traceback_msg
             else:
                 fn_call_result: FunctionCallResult = fn.invoke_fn_ser(
                     fn_name, input, init_value
@@ -153,11 +185,9 @@ def _run_function(
                 if fn_call_result.traceback_msg is not None:
                     print(fn_call_result.traceback_msg, file=sys.stderr)
                     has_failed = True
-                    exception_msg = fn_call_result.traceback_msg
-        except Exception as e:
+        except Exception:
             print(traceback.format_exc(), file=sys.stderr)
             has_failed = True
-            exception_msg = str(e)
 
     # WARNING - IF THIS FAILS, WE WILL NOT BE ABLE TO RECOVER
     # ANY LOGS
@@ -165,7 +195,6 @@ def _run_function(
         return FunctionOutput(
             fn_outputs=None,
             router_output=None,
-            exception=exception_msg,
             stdout=stdout_capture.getvalue(),
             stderr=stderr_capture.getvalue(),
             reducer=is_reducer,

indexify/executor/image_dependency_installer.py

@@ -7,7 +7,6 @@ from rich.theme import Theme
 
 from indexify.functions_sdk.image import ImageInformation
 
-
 custom_theme = Theme(
     {
         "info": "cyan",
@@ -20,14 +19,19 @@ custom_theme = Theme(
 console = Console(theme=custom_theme)
 
 
-def _record_image_name(name: str):
+def _record_image_name(name: str, version: int):
     dir_path = os.path.expanduser("~/.indexify/")
+
     file_path = os.path.expanduser("~/.indexify/image_name")
     os.makedirs(dir_path, exist_ok=True)
-
     with open(file_path, "w") as file:
         file.write(name)
 
+    file_path = os.path.expanduser("~/.indexify/image_version")
+    os.makedirs(dir_path, exist_ok=True)
+    with open(file_path, "w") as file:
+        file.write(str(version))
+
 
 def _install_dependencies(run_str: str):
     # Throw error to the caller if these subprocesses fail.
@@ -36,7 +40,9 @@ def _install_dependencies(run_str: str):
         raise Exception(f"Unable to install dep `{run_str}`")
 
 
-def executor_image_builder(image_info: ImageInformation, name_alias: str):
+def executor_image_builder(
+    image_info: ImageInformation, name_alias: str, image_version: int
+):
     console.print(Text("Attempting Executor Bootstrap.", style="red bold"))
 
     run_strs = image_info.run_strs
@@ -48,5 +54,11 @@ def executor_image_builder(image_info: ImageInformation, name_alias: str):
 
     console.print(Text("Install dependencies done.", style="red bold"))
 
-    console.print(Text(f"Recording image name {name_alias}", style="red bold"))
-    _record_image_name(name_alias)
+    console.print(
+        Text(
+            f"Recording image name {name_alias} and version {image_version}",
+            style="red bold",
+        )
+    )
+
+    _record_image_name(name_alias, image_version)

indexify/executor/runtime_probes.py

@@ -6,10 +6,12 @@ from typing import Any, Dict, Tuple
 from pydantic import BaseModel
 
 DEFAULT_EXECUTOR = "tensorlake/indexify-executor-default"
+DEFAULT_VERSION = 1
 
 
 class ProbeInfo(BaseModel):
     image_name: str
+    image_version: int
     python_major_version: int
     labels: Dict[str, Any] = {}
     is_default_executor: bool
@@ -18,6 +20,7 @@ class ProbeInfo(BaseModel):
 class RuntimeProbes:
     def __init__(self) -> None:
         self._image_name = self._read_image_name()
+        self._image_version = self._read_image_version()
         self._os_name = platform.system()
         self._architecture = platform.machine()
         (
@@ -32,6 +35,13 @@ class RuntimeProbes:
                 return file.read().strip()
         return DEFAULT_EXECUTOR
 
+    def _read_image_version(self) -> int:
+        file_path = os.path.expanduser("~/.indexify/image_version")
+        if os.path.exists(file_path):
+            with open(file_path, "r") as file:
+                return int(file.read().strip())
+        return DEFAULT_VERSION
+
     def _get_python_version(self) -> Tuple[int, int]:
         version_info = sys.version_info
         return version_info.major, version_info.minor
@@ -50,6 +60,7 @@ class RuntimeProbes:
 
         return ProbeInfo(
             image_name=self._image_name,
+            image_version=self._image_version,
             python_major_version=self._python_version_major,
             labels=labels,
             is_default_executor=self._is_default_executor(),

indexify/executor/task_reporter.py

@@ -32,22 +32,11 @@ class TaskReporter:
             print(
                 f"[bold]task-reporter[/bold] uploading output of size: {len(output.payload)} bytes"
             )
-            output_bytes = MsgPackSerializer.serialize(output) 
+            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)}"

indexify/executor/task_store.py

@@ -14,7 +14,6 @@ class CompletedTask(BaseModel):
     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

indexify/functions_sdk/data_objects.py

@@ -1,4 +1,4 @@
-from typing import Any, Dict, List, Optional, Union
+from typing import Any, Dict, List, Optional, Union, Literal
 
 from pydantic import BaseModel, Json
 
@@ -17,13 +17,12 @@ class RouterOutput(BaseModel):
 class IndexifyData(BaseModel):
     id: Optional[str] = None
     payload: bytes
-    payload_encoding: str = "cloudpickle"
+    encoder: Literal["cloudpickle", "json"] = "cloudpickle"
 
 
 class FunctionWorkerOutput(BaseModel):
     fn_outputs: Optional[List[IndexifyData]]
     router_output: Optional[RouterOutput]
-    exception: Optional[str]
     stdout: Optional[str]
     stderr: Optional[str]
     reducer: bool = False

indexify/functions_sdk/graph.py

@@ -33,6 +33,7 @@ from .indexify_functions import (
     IndexifyFunction,
     IndexifyFunctionWrapper,
     IndexifyRouter,
+    GraphInvocationContext,
 )
 from .local_cache import CacheAwareFunctionWrapper
 from .object_serializer import get_serializer
@@ -96,9 +97,9 @@ class Graph:
             return self
 
         if issubclass(indexify_fn, IndexifyFunction) and indexify_fn.accumulate:
-            self.accumulator_zero_values[indexify_fn.name] = (
-                indexify_fn.accumulate().model_dump()
-            )
+            self.accumulator_zero_values[
+                indexify_fn.name
+            ] = indexify_fn.accumulate().model_dump()
 
         self.nodes[indexify_fn.name] = indexify_fn
         return self
@@ -159,6 +160,7 @@ class Graph:
             reducer=start_node.accumulate is not None,
             image_name=start_node.image._image_name,
             image_information=start_node.image.to_image_information(),
+            payload_encoder=start_node.encoder
         )
         metadata_edges = self.edges.copy()
         metadata_nodes = {}
@@ -170,7 +172,7 @@ class Graph:
                         description=node.description or "",
                         source_fn=node_name,
                         target_fns=self.routers[node_name],
-                        payload_encoder=node.payload_encoder,
+                        payload_encoder=node.encoder,
                         image_name=node.image._image_name,
                         image_information=node.image.to_image_information(),
                     )
@@ -184,6 +186,7 @@ class Graph:
                         reducer=node.accumulate is not None,
                         image_name=node.image._image_name,
                         image_information=node.image.to_image_information(),
+                        encoder=node.encoder,
                     )
                 )
 
@@ -201,20 +204,26 @@ class Graph:
 
     def run(self, block_until_done: bool = False, **kwargs) -> str:
         start_node = self.nodes[self._start_node]
-        serializer = get_serializer(start_node.payload_encoder)
-        input = IndexifyData(id=generate(), payload=serializer.serialize(kwargs))
+        serializer = get_serializer(start_node.encoder)
+        input = IndexifyData(id=generate(), payload=serializer.serialize(kwargs), encoder=start_node.encoder)
         print(f"[bold] Invoking {self._start_node}[/bold]")
         outputs = defaultdict(list)
         self._accumulator_values[input.id] = {}
         for k, v in self.accumulator_zero_values.items():
             node = self.nodes[k]
-            serializer = get_serializer(node.payload_encoder)
+            serializer = get_serializer(node.encoder)
             self._accumulator_values[input.id] = {
-                k: IndexifyData(payload=serializer.serialize(v))
+                k: IndexifyData(payload=serializer.serialize(v), encoder=node.encoder)
             }
         self._results[input.id] = outputs
         enable_cache = kwargs.get("enable_cache", True)
-        self._run(input, outputs, enable_cache)
+        ctx = GraphInvocationContext(
+            invocation_id=input.id,
+            graph_name=self.name,
+            graph_version="1",
+            indexify_client=None,
+        )
+        self._run(input, outputs, enable_cache, ctx)
         return input.id
 
     def _run(
@@ -222,6 +231,7 @@ class Graph:
         initial_input: IndexifyData,
         outputs: Dict[str, List[bytes]],
         enable_cache: bool,
+        ctx: GraphInvocationContext,
     ):
         accumulator_values = self._accumulator_values[initial_input.id]
         queue = deque([(self._start_node, initial_input)])
@@ -229,7 +239,7 @@ class Graph:
             node_name, input = queue.popleft()
             node = self.nodes[node_name]
             function_outputs: FunctionCallResult = IndexifyFunctionWrapper(
-                node
+                node, context=ctx
             ).invoke_fn_ser(node_name, input, accumulator_values.get(node_name, None))
             if function_outputs.traceback_msg is not None:
                 print(function_outputs.traceback_msg)
@@ -281,7 +291,7 @@ class Graph:
             raise ValueError(f"no results found for fn {fn_name} on graph {self.name}")
         fn = self.nodes[fn_name]
         fn_model = self.get_function(fn_name).get_output_model()
-        serializer = get_serializer(fn.payload_encoder)
+        serializer = get_serializer(fn.encoder)
         outputs = []
         for result in results[fn_name]:
             payload_dict = serializer.deserialize(result.payload)

indexify/functions_sdk/graph_definition.py

@@ -48,4 +48,4 @@ class ComputeGraphMetadata(BaseModel):
     runtime_information: RuntimeInformation
 
     def get_input_payload_serializer(self):
-        return get_serializer(self.start_node.compute_fn.payload_encoder)
+        return get_serializer(self.start_node.compute_fn.encoder)

indexify/functions_sdk/indexify_functions.py

@@ -2,7 +2,6 @@ import inspect
 import re
 import sys
 import traceback
-from abc import ABC, abstractmethod
 from functools import update_wrapper
 from typing import (
     Any,
@@ -18,12 +17,35 @@ from typing import (
 )
 
 import msgpack
-from pydantic import BaseModel
+from pydantic import BaseModel, Field, PrivateAttr, model_validator
 from typing_extensions import get_type_hints
 
-from .data_objects import IndexifyData, RouterOutput
+from .data_objects import IndexifyData
 from .image import DEFAULT_IMAGE_3_10, Image
-from .object_serializer import CloudPickleSerializer, get_serializer
+from .object_serializer import get_serializer
+
+
+class GraphInvocationContext(BaseModel):
+    invocation_id: str
+    graph_name: str
+    graph_version: str
+    indexify_client: Optional[Any] = Field(default=None)  # avoids circular import
+    _local_state: Dict[str, Any] = PrivateAttr(default_factory=dict)
+
+    def set_state_key(self, key: str, value: Any) -> None:
+        if self.indexify_client is None:
+            self._local_state[key] = value
+            return
+        self.indexify_client.set_state_key(
+            self.graph_name, self.invocation_id, key, value
+        )
+
+    def get_state_key(self, key: str) -> Any:
+        if self.indexify_client is None:
+            return self._local_state.get(key)
+        return self.indexify_client.get_state_key(
+            self.graph_name, self.invocation_id, key
+        )
 
 
 def format_filtered_traceback(exc_info=None):
@@ -97,7 +119,7 @@ class IndexifyFunction:
     image: Optional[Image] = DEFAULT_IMAGE_3_10
     placement_constraints: List[PlacementConstraints] = []
     accumulate: Optional[Type[Any]] = None
-    payload_encoder: Optional[str] = "cloudpickle"
+    encoder: Optional[str] = "cloudpickle"
 
     def run(self, *args, **kwargs) -> Union[List[Any], Any]:
         pass
@@ -109,7 +131,7 @@ class IndexifyFunction:
 
     @classmethod
     def deserialize_output(cls, output: IndexifyData) -> Any:
-        serializer = get_serializer(cls.payload_encoder)
+        serializer = get_serializer(cls.encoder)
         return serializer.deserialize(output.payload)
 
 
@@ -118,7 +140,7 @@ class IndexifyRouter:
     description: str = ""
     image: Optional[Image] = DEFAULT_IMAGE_3_10
     placement_constraints: List[PlacementConstraints] = []
-    payload_encoder: Optional[str] = "cloudpickle"
+    encoder: Optional[str] = "cloudpickle"
 
     def run(self, *args, **kwargs) -> Optional[List[IndexifyFunction]]:
         pass
@@ -129,7 +151,7 @@ def indexify_router(
     description: Optional[str] = "",
     image: Optional[Image] = DEFAULT_IMAGE_3_10,
     placement_constraints: List[PlacementConstraints] = [],
-    payload_encoder: Optional[str] = "cloudpickle",
+    encoder: Optional[str] = "cloudpickle",
 ):
     def construct(fn):
         args = locals().copy()
@@ -152,7 +174,7 @@ def indexify_router(
                 setattr(IndexifyRo, key, value)
 
         IndexifyRo.image = image
-        IndexifyRo.payload_encoder = payload_encoder
+        IndexifyRo.encoder = encoder
         return IndexifyRo
 
     return construct
@@ -163,7 +185,7 @@ def indexify_function(
     description: Optional[str] = "",
     image: Optional[Image] = DEFAULT_IMAGE_3_10,
     accumulate: Optional[Type[BaseModel]] = None,
-    payload_encoder: Optional[str] = "cloudpickle",
+    encoder: Optional[str] = "cloudpickle",
     placement_constraints: List[PlacementConstraints] = [],
 ):
     def construct(fn):
@@ -185,10 +207,9 @@ def indexify_function(
         for key, value in args.items():
             if key != "fn" and key != "self":
                 setattr(IndexifyFn, key, value)
-
         IndexifyFn.image = image
         IndexifyFn.accumulate = accumulate
-        IndexifyFn.payload_encoder = payload_encoder
+        IndexifyFn.encoder = encoder
         return IndexifyFn
 
     return construct
@@ -205,10 +226,15 @@ class RouterCallResult(BaseModel):
 
 
 class IndexifyFunctionWrapper:
-    def __init__(self, indexify_function: Union[IndexifyFunction, IndexifyRouter]):
-        self.indexify_function: Union[IndexifyFunction, IndexifyRouter] = (
-            indexify_function()
-        )
+    def __init__(
+        self,
+        indexify_function: Union[IndexifyFunction, IndexifyRouter],
+        context: GraphInvocationContext,
+    ):
+        self.indexify_function: Union[
+            IndexifyFunction, IndexifyRouter
+        ] = indexify_function()
+        self.indexify_function._ctx = context
 
     def get_output_model(self) -> Any:
         if not isinstance(self.indexify_function, IndexifyFunction):
@@ -276,7 +302,7 @@ class IndexifyFunctionWrapper:
         self, name: str, input: IndexifyData, acc: Optional[Any] = None
     ) -> FunctionCallResult:
         input = self.deserialize_input(name, input)
-        serializer = get_serializer(self.indexify_function.payload_encoder)
+        serializer = get_serializer(self.indexify_function.encoder)
         if acc is not None:
             acc = self.indexify_function.accumulate.model_validate(
                 serializer.deserialize(acc.payload)
@@ -287,7 +313,7 @@ class IndexifyFunctionWrapper:
             )
         outputs, err = self.run_fn(input, acc=acc)
         ser_outputs = [
-            IndexifyData(payload=serializer.serialize(output)) for output in outputs
+            IndexifyData(payload=serializer.serialize(output), encoder=self.indexify_function.encoder) for output in outputs
         ]
         return FunctionCallResult(ser_outputs=ser_outputs, traceback_msg=err)
 
@@ -297,28 +323,17 @@ class IndexifyFunctionWrapper:
         return RouterCallResult(edges=edges, traceback_msg=err)
 
     def deserialize_input(self, compute_fn: str, indexify_data: IndexifyData) -> Any:
-        if self.indexify_function.payload_encoder == "cloudpickle":
-            return CloudPickleSerializer.deserialize(indexify_data.payload)
-        payload = msgpack.unpackb(indexify_data.payload)
-        signature = inspect.signature(self.indexify_function.run)
-        arg_types = {}
-        for name, param in signature.parameters.items():
-            if (
-                param.annotation != inspect.Parameter.empty
-                and param.annotation != getattr(compute_fn, "accumulate", None)
-            ):
-                arg_types[name] = param.annotation
-        if len(arg_types) > 1:
-            raise ValueError(
-                f"Compute function {compute_fn} has multiple arguments, but only one is supported"
-            )
-        elif len(arg_types) == 0:
-            raise ValueError(f"Compute function {compute_fn} has no arguments")
-        arg_name, arg_type = next(iter(arg_types.items()))
-        if arg_type is None:
-            raise ValueError(f"Argument {arg_name} has no type annotation")
-        if is_pydantic_model_from_annotation(arg_type):
-            if len(payload.keys()) == 1 and isinstance(list(payload.values())[0], dict):
-                payload = list(payload.values())[0]
-            return arg_type.model_validate(payload)
-        return payload
+        encoder = indexify_data.encoder
+        payload = indexify_data.payload
+        serializer = get_serializer(encoder)
+        return serializer.deserialize(payload)
+
+def get_ctx() -> GraphInvocationContext:
+    frame = inspect.currentframe()
+    caller_frame = frame.f_back.f_back
+    function_instance = caller_frame.f_locals["self"]
+    del frame
+    del caller_frame
+    if isinstance(function_instance, IndexifyFunctionWrapper):
+        return function_instance.indexify_function._ctx
+    return function_instance._ctx

indexify/functions_sdk/object_serializer.py

@@ -1,5 +1,6 @@
 from typing import Any, List
 
+import jsonpickle
 import cloudpickle
 import msgpack
 from pydantic import BaseModel
@@ -12,8 +13,26 @@ def get_serializer(serializer_type: str) -> Any:
         return CloudPickleSerializer()
     elif serializer_type == "msgpack":
         return MsgPackSerializer()
-    else:
-        raise ValueError(f"Unknown serializer type: {serializer_type}")
+    elif serializer_type == "json":
+        return JsonSerializer()
+    raise ValueError(f"Unknown serializer type: {serializer_type}")
+
+class JsonSerializer:
+    @staticmethod
+    def serialize(data: Any) -> str:
+        return jsonpickle.encode(data)
+
+    @staticmethod
+    def deserialize(data: str) -> Any:
+        return jsonpickle.decode(data)
+
+    @staticmethod
+    def serialize_list(data: List[Any]) -> str:
+        return jsonpickle.encode(data)
+
+    @staticmethod
+    def deserialize_list(data: str) -> List[Any]:
+        return jsonpickle.decode(data)
 
 
 class CloudPickleSerializer:

indexify/functions_sdk/pipeline.py

@@ -28,3 +28,6 @@ class Pipeline:
     def run(self, **kwargs):
         invocation_id = self._graph.run(**kwargs)
         return invocation_id
+
+    def output(self, invocation_id: str, function_name: str):
+        return self._graph.output(invocation_id, function_name)

indexify/http_client.py

@@ -14,8 +14,7 @@ from indexify.error import ApiException, GraphStillProcessing
 from indexify.functions_sdk.data_objects import IndexifyData
 from indexify.functions_sdk.graph import ComputeGraphMetadata, Graph
 from indexify.functions_sdk.indexify_functions import IndexifyFunction
-from indexify.settings import DEFAULT_SERVICE_URL, DEFAULT_SERVICE_URL_HTTPS
-
+from indexify.settings import DEFAULT_SERVICE_URL
 
 class InvocationEventPayload(BaseModel):
     invocation_id: str
@@ -47,6 +46,7 @@ class IndexifyClient:
         service_url: str = DEFAULT_SERVICE_URL,
         config_path: Optional[str] = None,
         namespace: str = "default",
+        api_key: Optional[str] = None,
         **kwargs,
     ):
         if os.environ.get("INDEXIFY_URL"):
@@ -74,6 +74,10 @@ class IndexifyClient:
         self._timeout = kwargs.get("timeout")
         self._graphs: Dict[str, Graph] = {}
         self._fns: Dict[str, IndexifyFunction] = {}
+        self._api_key = api_key
+        if not self._api_key:
+            print("API key not provided. Trying to fetch from environment TENSORLAKE_API_KEY variable")
+            self._api_key = os.getenv("TENSORLAKE_API_KEY")
 
     def _request(self, method: str, **kwargs) -> httpx.Response:
         try:
@@ -87,10 +91,9 @@ class IndexifyClient:
                 raise ApiException(response.text)
         except httpx.ConnectError:
             message = (
-                f"Make sure the server is running and accesible at {self._service_url}"
+                f"Make sure the server is running and accessible at {self._service_url}"
             )
             ex = ApiException(message=message)
-            print(ex)
             raise ex
         return response
 
@@ -100,7 +103,7 @@ class IndexifyClient:
         cert_path: str,
         key_path: str,
         ca_bundle_path: Optional[str] = None,
-        service_url: str = DEFAULT_SERVICE_URL_HTTPS,
+        service_url: str = DEFAULT_SERVICE_URL,
         *args,
         **kwargs,
     ) -> "IndexifyClient":
@@ -140,17 +143,25 @@ class IndexifyClient:
             verify=verify_option,
         )
         return client
+    
+    def _add_api_key(self, kwargs):
+        if self._api_key:
+            kwargs["headers"] = {"Authorization": f"Bearer {self._api_key}"}
 
     def _get(self, endpoint: str, **kwargs) -> httpx.Response:
+        self._add_api_key(kwargs)
         return self._request("GET", url=f"{self._service_url}/{endpoint}", **kwargs)
 
     def _post(self, endpoint: str, **kwargs) -> httpx.Response:
+        self._add_api_key(kwargs)
         return self._request("POST", url=f"{self._service_url}/{endpoint}", **kwargs)
 
     def _put(self, endpoint: str, **kwargs) -> httpx.Response:
+        self._add_api_key(kwargs)
         return self._request("PUT", url=f"{self._service_url}/{endpoint}", **kwargs)
 
     def _delete(self, endpoint: str, **kwargs) -> httpx.Response:
+        self._add_api_key(kwargs)
         return self._request("DELETE", url=f"{self._service_url}/{endpoint}", **kwargs)
 
     def _close(self):
@@ -198,6 +209,23 @@ class IndexifyClient:
             namespaces.append(item["name"])
         return namespaces
 
+    def set_state_key(
+        self, compute_graph: str, invocation_id: str, key: str, value: Json
+    ) -> None:
+        response = self._post(
+                f"internal/namespaces/{self.namespace}/compute_graphs/{compute_graph}/invocations/{invocation_id}/ctx",
+                json={"key": key, "value": value},
+            )
+        response.raise_for_status()
+
+    def get_state_key(self, compute_graph: str, invocation_id: str, key: str) -> Json:
+        response = self._get(
+                f"internal/namespaces/{self.namespace}/compute_graphs/{compute_graph}/invocations/{invocation_id}/ctx",
+                json={"key": key},
+            )
+        response.raise_for_status()
+        return response.json().get("value")
+
     @classmethod
     def new_namespace(
         cls, namespace: str, server_addr: Optional[str] = "http://localhost:8900"
@@ -242,15 +270,18 @@ class IndexifyClient:
     ) -> str:
         ser_input = cloudpickle.dumps(kwargs)
         params = {"block_until_finish": block_until_done}
+        kwargs = {"headers": {"Content-Type": "application/cbor"}, "data": ser_input, "params":params}
+        self._add_api_key(kwargs)
         with httpx.Client() as client:
             with connect_sse(
                 client,
                 "POST",
                 f"{self.service_url}/namespaces/{self.namespace}/compute_graphs/{graph}/invoke_object",
-                headers={"Content-Type": "application/cbor"},
-                data=ser_input,
-                params=params,
+                **kwargs,
             ) as event_source:
+                if not event_source.response.is_success:
+                    resp = event_source.response.read().decode("utf-8")
+                    raise Exception(f"failed to invoke graph: {resp}")
                 for sse in event_source.iter_sse():
                     obj = json.loads(sse.data)
                     for k, v in obj.items():

indexify/remote_graph.py

@@ -84,7 +84,12 @@ class RemoteGraph:
         return cls(name=g.name, server_url=server_url, client=client)
 
     @classmethod
-    def by_name(cls, name: str, server_url: Optional[str] = DEFAULT_SERVICE_URL, client: Optional[IndexifyClient] = None):
+    def by_name(
+        cls,
+        name: str,
+        server_url: Optional[str] = DEFAULT_SERVICE_URL,
+        client: Optional[IndexifyClient] = None,
+    ):
         """
         Create a handle to call a RemoteGraph by name.
 
@@ -104,7 +109,7 @@ class RemoteGraph:
     ) -> List[Any]:
         """
         Returns the extracted objects by a graph for an ingested object.
- 
+
         - If the extractor name is provided, only the objects extracted by that extractor are returned.
         - If the extractor name is not provided, all the extracted objects are returned for the input object.
 

indexify/settings.py

@@ -1,2 +1 @@
-DEFAULT_SERVICE_URL = "http://localhost:8900"
-DEFAULT_SERVICE_URL_HTTPS = "https://localhost:8900"
+DEFAULT_SERVICE_URL = "https://api.tensorlake.ai"

{indexify-0.2.22.dist-info → indexify-0.2.24.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: indexify
-Version: 0.2.22
+Version: 0.2.24
 Summary: Python Client for Indexify
 Home-page: https://github.com/tensorlakeai/indexify
 License: Apache 2.0
@@ -18,6 +18,7 @@ Requires-Dist: cloudpickle (>=3.1.0,<4.0.0)
 Requires-Dist: docker (>=7.1.0,<8.0.0)
 Requires-Dist: httpx-sse (>=0.4.0,<0.5.0)
 Requires-Dist: httpx[http2] (>=0,<1)
+Requires-Dist: jsonpickle (>=3.3.0,<4.0.0)
 Requires-Dist: msgpack (>=1.1.0,<2.0.0)
 Requires-Dist: nanoid (>=2.0.0,<3.0.0)
 Requires-Dist: pydantic (>=2.9.2,<3.0.0)

indexify-0.2.24.dist-info/RECORD

@@ -0,0 +1,34 @@
+indexify/__init__.py,sha256=P0mvM8sbkeS2CjYzRYyzb42CnXGhyJXdz4FdmTBMSWM,697
+indexify/cli.py,sha256=TBycogrXcTK7q8IIcJxcpPPoZYve_rp6FYipWbkiTUI,8528
+indexify/data_loaders/__init__.py,sha256=Y5NEuseTcYAICRiweYw5wBQ2m2YplbsY21I7df-rdi4,1339
+indexify/data_loaders/local_directory_loader.py,sha256=fCrgj5drnW71ZUdDDvcB1-VJjIs1w6Q8sEW0HSGSAiA,1247
+indexify/data_loaders/url_loader.py,sha256=32SERljcq1Xsi4RdLz2dgyk2TER5pQPTtXl3gUzwHbY,1533
+indexify/error.py,sha256=qAWr8R6AxPkjsxHSzXTc8zqYnNO_AjOqqYEPsQvF1Zs,238
+indexify/executor/agent.py,sha256=3BpvTlmLIjBjrQIaTrQlMy4TLu3wJL_2Nwb6wZJGM8k,18692
+indexify/executor/api_objects.py,sha256=mvmwGbK4paJNQGFvbtNHMPpiH_LpVhrlRnCcrqS6HOQ,859
+indexify/executor/downloader.py,sha256=3mEDdluTzspsLGAZtFHZOVuyKOzT3CSema2kIK6Z1yU,4005
+indexify/executor/executor_tasks.py,sha256=A0UIEZ5VaB6zSkFQG81UmTW0E57MTYhGlaXuAbRV8lQ,1884
+indexify/executor/function_worker.py,sha256=bvmHpcDOcEh4Hy7IK0dRWu79bU96rs_j8o3utUR_5-A,6416
+indexify/executor/image_dependency_installer.py,sha256=ct8GmzgkaPi6NAblk68IJJWo5MecIUubELotmSrgoRQ,1759
+indexify/executor/indexify_executor.py,sha256=2Ut_VX-Su_lm4b4aEROyRJ3gXx-uFHA-V7EN0sWiARE,771
+indexify/executor/runtime_probes.py,sha256=mjw2_mGQ622wRT_39WPGGgPEZQTgtrf3-ICcUUZOeyg,2126
+indexify/executor/task_reporter.py,sha256=7M2fDzLDkSNOSr42EBtkIbGvBzeZlTCqxQuDfdoFT4Y,3349
+indexify/executor/task_store.py,sha256=u48FdRKAh_KH7WOMQOArdOY5CawlyW5MJx8V0W79JM0,3951
+indexify/functions_sdk/data_objects.py,sha256=BJjIQJ_wSyRhQI50Y9_y_zLwCLkDiVC2CrAck-UZyvw,842
+indexify/functions_sdk/graph.py,sha256=ZSOLuDt2PZeFecxm7oLfML57oQaFdDN0hmB4gr_pQTk,11518
+indexify/functions_sdk/graph_definition.py,sha256=s8NbqZRaVw1mhGvaaUZ2cJ0NitdlbLj1tnkKjsc71U4,1309
+indexify/functions_sdk/graph_validation.py,sha256=XLHiC9PAtZungJLysU3hIUOPNDkO5TXUDZ_jiZ0H4hg,2508
+indexify/functions_sdk/image.py,sha256=QK0H6KxLWriB_z4M0kunKzzHdHxYLWL670RPYgYuf_8,1762
+indexify/functions_sdk/indexify_functions.py,sha256=6KbX2c75XNbOToCIuYmamWLJE5RIjUABsHkvARQXIsI,11162
+indexify/functions_sdk/local_cache.py,sha256=cNWF67zbhbTJe3g86hyLBy3Rqzs6dNvp2SjLazGZWvw,1348
+indexify/functions_sdk/object_serializer.py,sha256=Zy9IVqAvylCdKgNE8IpuwCd-HDN-IP7fOcWd8RQ0by8,2160
+indexify/functions_sdk/pipeline.py,sha256=KmxZE8eBFAQ4bbEcYURXXR26HSyoAT3O6iu9H38-OXE,974
+indexify/http_client.py,sha256=fg5gGIdAhz9ZM9L4XmP9-9X8LJp5E2sMnkXMSQL9EaQ,15349
+indexify/remote_graph.py,sha256=_04bdraDaeEmkvEcsxIZ8M37BCAjfbfQccZpAgnZJ2c,4435
+indexify/remote_pipeline.py,sha256=FW7IAv3r24OOpiqlprw3kuFrpdkqi6Ic4_tT26FThjA,761
+indexify/settings.py,sha256=Ny59mzYI4gbXoK8hjx66a_men6ndbd1J1zCTcKOoyzg,50
+indexify-0.2.24.dist-info/LICENSE.txt,sha256=xx0jnfkXJvxRnG63LTGOxlggYnIysveWIZ6H3PNdCrQ,11357
+indexify-0.2.24.dist-info/METADATA,sha256=yrUA8bcrVZcWXD5oy-ZF06-9A3zHCoU0hhComMhRRbM,6242
+indexify-0.2.24.dist-info/WHEEL,sha256=Nq82e9rUAnEjt98J6MlVmMCZb-t9cYE2Ir1kpBmnWfs,88
+indexify-0.2.24.dist-info/entry_points.txt,sha256=Pih7WV-XMpAzI5dEvROcpLr-ybVhd9Y-AtuzBKUdcDs,49
+indexify-0.2.24.dist-info/RECORD,,

indexify-0.2.22.dist-info/RECORD

@@ -1,34 +0,0 @@
-indexify/__init__.py,sha256=xYjdqZQ9CEtV-s4PJhj_6Vx_dfg_qknWHkLqOSVom8g,623
-indexify/cli.py,sha256=sR-9wc3J_46TsOV424VD6TwKZkEg9MHiflvAJB1VyeY,8420
-indexify/data_loaders/__init__.py,sha256=Y5NEuseTcYAICRiweYw5wBQ2m2YplbsY21I7df-rdi4,1339
-indexify/data_loaders/local_directory_loader.py,sha256=fCrgj5drnW71ZUdDDvcB1-VJjIs1w6Q8sEW0HSGSAiA,1247
-indexify/data_loaders/url_loader.py,sha256=32SERljcq1Xsi4RdLz2dgyk2TER5pQPTtXl3gUzwHbY,1533
-indexify/error.py,sha256=mBEz_fj2sdtTQol-VfwvD8Syu7x9uvYSy2nne4Gr2GI,237
-indexify/executor/agent.py,sha256=nSJ9vIWqY74pL5cEy7Rvb_xDatqtgEmxT4sJ65z3Rb0,17984
-indexify/executor/api_objects.py,sha256=QMlHbcL-ZOAMVD71XOJqD_ZIy-NjrUf1DAvi0hFYVg0,836
-indexify/executor/downloader.py,sha256=3mEDdluTzspsLGAZtFHZOVuyKOzT3CSema2kIK6Z1yU,4005
-indexify/executor/executor_tasks.py,sha256=gAZ2pvza1YwGlaR1o_tJW4SXtdCgK7sLJgp4W7rOjR0,1834
-indexify/executor/function_worker.py,sha256=pCGn13rg4dEykzmwYNyCTkewrpyQXQR1cH6n2Hx5Lfc,5813
-indexify/executor/image_dependency_installer.py,sha256=ditkQQ8rMrtQgZZ7Nw5wDn4I0d-qk8Q2irL-cU-4gOY,1449
-indexify/executor/indexify_executor.py,sha256=2Ut_VX-Su_lm4b4aEROyRJ3gXx-uFHA-V7EN0sWiARE,771
-indexify/executor/runtime_probes.py,sha256=JY0FoxtlQ9sgsE8gBKWM5h3R1TWkYENhNF0HR2KkV4Q,1704
-indexify/executor/task_reporter.py,sha256=0ujOJrN7aW9PiVH0gr1rGqadfZE9yfV318Ll3eahBng,3721
-indexify/executor/task_store.py,sha256=q8s2gImsFffWeXQR0mk1Xlo1Aj_2GfclNPjQ2EA_YBo,3984
-indexify/functions_sdk/data_objects.py,sha256=CQZMzYiV7l6dyzFkYquWQHqdte6JnC7XA3i2ZyvPvgQ,844
-indexify/functions_sdk/graph.py,sha256=L8jtuSGaUk05JgT_cQX-n_1H9t84lcPa1v4XXlpuJa0,11143
-indexify/functions_sdk/graph_definition.py,sha256=fwv63wkuKeVQxJQk9ofu8ZL3B0NzJzF9bJhFdVZZvMQ,1317
-indexify/functions_sdk/graph_validation.py,sha256=XLHiC9PAtZungJLysU3hIUOPNDkO5TXUDZ_jiZ0H4hg,2508
-indexify/functions_sdk/image.py,sha256=QK0H6KxLWriB_z4M0kunKzzHdHxYLWL670RPYgYuf_8,1762
-indexify/functions_sdk/indexify_functions.py,sha256=KLeRQY5IcrVp5T3K4yQ8Uu_CLL-Cj7uOG0L7y6uJnfU,11109
-indexify/functions_sdk/local_cache.py,sha256=cNWF67zbhbTJe3g86hyLBy3Rqzs6dNvp2SjLazGZWvw,1348
-indexify/functions_sdk/object_serializer.py,sha256=Zz4GobW3ZamBBtFDF76QxU3TP6oJNdWnhsfKd0OUFoc,1660
-indexify/functions_sdk/pipeline.py,sha256=7hDatRK-SCHYvttf2Vj5YFyiJEVU0OOXEZBOIQenSmk,847
-indexify/http_client.py,sha256=pFJ4i40_GgzGJP9ih_iCHNUJQlSsWUiCLBMT3aIegOg,13887
-indexify/remote_graph.py,sha256=NOotayhw_x4-mAw1c3ooNe7gFE_q3QJ9qaO0t0Sb8sE,4397
-indexify/remote_pipeline.py,sha256=FW7IAv3r24OOpiqlprw3kuFrpdkqi6Ic4_tT26FThjA,761
-indexify/settings.py,sha256=LSaWZ0ADIVmUv6o6dHWRC3-Ry5uLbCw2sBSg1e_U7UM,99
-indexify-0.2.22.dist-info/LICENSE.txt,sha256=xx0jnfkXJvxRnG63LTGOxlggYnIysveWIZ6H3PNdCrQ,11357
-indexify-0.2.22.dist-info/METADATA,sha256=PTlq-SzEMf1hcd09SwyB7aBF4BvsQgVPyPK8PQoGrCU,6199
-indexify-0.2.22.dist-info/WHEEL,sha256=Nq82e9rUAnEjt98J6MlVmMCZb-t9cYE2Ir1kpBmnWfs,88
-indexify-0.2.22.dist-info/entry_points.txt,sha256=Pih7WV-XMpAzI5dEvROcpLr-ybVhd9Y-AtuzBKUdcDs,49
-indexify-0.2.22.dist-info/RECORD,,