indexify 0.2.22__tar.gz → 0.2.23__tar.gz

{indexify-0.2.22 → indexify-0.2.23}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: indexify
-Version: 0.2.22
+Version: 0.2.23
 Summary: Python Client for Indexify
 Home-page: https://github.com/tensorlakeai/indexify
 License: Apache 2.0

{indexify-0.2.22 → indexify-0.2.23}/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-0.2.22 → indexify-0.2.23}/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-0.2.22 → indexify-0.2.23}/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-0.2.22 → indexify-0.2.23}/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
@@ -22,8 +20,11 @@ from indexify.functions_sdk.data_objects import (
     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 +32,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 +60,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 +112,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 +201,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:
@@ -378,11 +384,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 +419,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-0.2.22 → indexify-0.2.23}/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-0.2.22 → indexify-0.2.23}/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-0.2.22 → indexify-0.2.23}/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,
 )
@@ -44,7 +46,13 @@ class FunctionOutput(BaseModel):
 
 
 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 +62,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,10 +94,19 @@ 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:
@@ -113,6 +140,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
@@ -131,7 +160,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 (

{indexify-0.2.22 → indexify-0.2.23}/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-0.2.22 → indexify-0.2.23}/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-0.2.22 → indexify-0.2.23}/indexify/executor/task_reporter.py

@@ -32,7 +32,7 @@ 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)))
             )

{indexify-0.2.22 → indexify-0.2.23}/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
@@ -214,7 +215,13 @@ class Graph:
             }
         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 +229,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 +237,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)

{indexify-0.2.22 → indexify-0.2.23}/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,14 +17,37 @@ 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
 
 
+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):
     """
     Format a traceback excluding indexify_functions.py lines.
@@ -205,10 +227,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):
@@ -322,3 +349,14 @@ class IndexifyFunctionWrapper:
                 payload = list(payload.values())[0]
             return arg_type.model_validate(payload)
         return 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-0.2.22 → indexify-0.2.23}/indexify/http_client.py

@@ -198,6 +198,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"
@@ -251,6 +268,9 @@ class IndexifyClient:
                 data=ser_input,
                 params=params,
             ) 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-0.2.22 → indexify-0.2.23}/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-0.2.22 → indexify-0.2.23}/pyproject.toml

@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "indexify"
-version = "0.2.22"
+version = "0.2.23"
 description = "Python Client for Indexify"
 authors = ["Tensorlake Inc. <support@tensorlake.ai>"]
 license = "Apache 2.0"