indexify 0.2.44__py3-none-any.whl → 0.2.45__py3-none-any.whl

indexify/function_executor/proto/function_executor_pb2.pyi

@@ -52,6 +52,76 @@ class InitializeResponse(_message.Message):
     success: bool
     def __init__(self, success: bool = ...) -> None: ...
 
+class SetInvocationStateRequest(_message.Message):
+    __slots__ = ("key", "value")
+    KEY_FIELD_NUMBER: _ClassVar[int]
+    VALUE_FIELD_NUMBER: _ClassVar[int]
+    key: str
+    value: SerializedObject
+    def __init__(
+        self,
+        key: _Optional[str] = ...,
+        value: _Optional[_Union[SerializedObject, _Mapping]] = ...,
+    ) -> None: ...
+
+class SetInvocationStateResponse(_message.Message):
+    __slots__ = ()
+    def __init__(self) -> None: ...
+
+class GetInvocationStateRequest(_message.Message):
+    __slots__ = ("key",)
+    KEY_FIELD_NUMBER: _ClassVar[int]
+    key: str
+    def __init__(self, key: _Optional[str] = ...) -> None: ...
+
+class GetInvocationStateResponse(_message.Message):
+    __slots__ = ("key", "value")
+    KEY_FIELD_NUMBER: _ClassVar[int]
+    VALUE_FIELD_NUMBER: _ClassVar[int]
+    key: str
+    value: SerializedObject
+    def __init__(
+        self,
+        key: _Optional[str] = ...,
+        value: _Optional[_Union[SerializedObject, _Mapping]] = ...,
+    ) -> None: ...
+
+class InvocationStateRequest(_message.Message):
+    __slots__ = ("request_id", "task_id", "set", "get")
+    REQUEST_ID_FIELD_NUMBER: _ClassVar[int]
+    TASK_ID_FIELD_NUMBER: _ClassVar[int]
+    SET_FIELD_NUMBER: _ClassVar[int]
+    GET_FIELD_NUMBER: _ClassVar[int]
+    request_id: str
+    task_id: str
+    set: SetInvocationStateRequest
+    get: GetInvocationStateRequest
+    def __init__(
+        self,
+        request_id: _Optional[str] = ...,
+        task_id: _Optional[str] = ...,
+        set: _Optional[_Union[SetInvocationStateRequest, _Mapping]] = ...,
+        get: _Optional[_Union[GetInvocationStateRequest, _Mapping]] = ...,
+    ) -> None: ...
+
+class InvocationStateResponse(_message.Message):
+    __slots__ = ("request_id", "success", "set", "get")
+    REQUEST_ID_FIELD_NUMBER: _ClassVar[int]
+    SUCCESS_FIELD_NUMBER: _ClassVar[int]
+    SET_FIELD_NUMBER: _ClassVar[int]
+    GET_FIELD_NUMBER: _ClassVar[int]
+    request_id: str
+    success: bool
+    set: SetInvocationStateResponse
+    get: GetInvocationStateResponse
+    def __init__(
+        self,
+        request_id: _Optional[str] = ...,
+        success: bool = ...,
+        set: _Optional[_Union[SetInvocationStateResponse, _Mapping]] = ...,
+        get: _Optional[_Union[GetInvocationStateResponse, _Mapping]] = ...,
+    ) -> None: ...
+
 class FunctionOutput(_message.Message):
     __slots__ = ("outputs",)
     OUTPUTS_FIELD_NUMBER: _ClassVar[int]

indexify/function_executor/proto/function_executor_pb2_grpc.py

@@ -46,6 +46,12 @@ class FunctionExecutorStub(object):
             response_deserializer=indexify_dot_function__executor_dot_proto_dot_function__executor__pb2.InitializeResponse.FromString,
             _registered_method=True,
         )
+        self.initialize_invocation_state_server = channel.stream_stream(
+            "/function_executor_service.FunctionExecutor/initialize_invocation_state_server",
+            request_serializer=indexify_dot_function__executor_dot_proto_dot_function__executor__pb2.InvocationStateResponse.SerializeToString,
+            response_deserializer=indexify_dot_function__executor_dot_proto_dot_function__executor__pb2.InvocationStateRequest.FromString,
+            _registered_method=True,
+        )
         self.run_task = channel.unary_unary(
             "/function_executor_service.FunctionExecutor/run_task",
             request_serializer=indexify_dot_function__executor_dot_proto_dot_function__executor__pb2.RunTaskRequest.SerializeToString,
@@ -67,6 +73,15 @@ class FunctionExecutorServicer(object):
         context.set_details("Method not implemented!")
         raise NotImplementedError("Method not implemented!")
 
+    def initialize_invocation_state_server(self, request_iterator, context):
+        """Initializes a server that sends requests to the client to perform actions on
+        a task's graph invocation state. This method is called only once per Function Executor
+        It should be called before calling RunTask for the function.
+        """
+        context.set_code(grpc.StatusCode.UNIMPLEMENTED)
+        context.set_details("Method not implemented!")
+        raise NotImplementedError("Method not implemented!")
+
     def run_task(self, request, context):
         """Executes the task defined in the request.
         Multiple tasks can be running in parallel.
@@ -83,6 +98,11 @@ def add_FunctionExecutorServicer_to_server(servicer, server):
             request_deserializer=indexify_dot_function__executor_dot_proto_dot_function__executor__pb2.InitializeRequest.FromString,
             response_serializer=indexify_dot_function__executor_dot_proto_dot_function__executor__pb2.InitializeResponse.SerializeToString,
         ),
+        "initialize_invocation_state_server": grpc.stream_stream_rpc_method_handler(
+            servicer.initialize_invocation_state_server,
+            request_deserializer=indexify_dot_function__executor_dot_proto_dot_function__executor__pb2.InvocationStateResponse.FromString,
+            response_serializer=indexify_dot_function__executor_dot_proto_dot_function__executor__pb2.InvocationStateRequest.SerializeToString,
+        ),
         "run_task": grpc.unary_unary_rpc_method_handler(
             servicer.run_task,
             request_deserializer=indexify_dot_function__executor_dot_proto_dot_function__executor__pb2.RunTaskRequest.FromString,
@@ -132,6 +152,36 @@ class FunctionExecutor(object):
             _registered_method=True,
         )
 
+    @staticmethod
+    def initialize_invocation_state_server(
+        request_iterator,
+        target,
+        options=(),
+        channel_credentials=None,
+        call_credentials=None,
+        insecure=False,
+        compression=None,
+        wait_for_ready=None,
+        timeout=None,
+        metadata=None,
+    ):
+        return grpc.experimental.stream_stream(
+            request_iterator,
+            target,
+            "/function_executor_service.FunctionExecutor/initialize_invocation_state_server",
+            indexify_dot_function__executor_dot_proto_dot_function__executor__pb2.InvocationStateResponse.SerializeToString,
+            indexify_dot_function__executor_dot_proto_dot_function__executor__pb2.InvocationStateRequest.FromString,
+            options,
+            channel_credentials,
+            insecure,
+            call_credentials,
+            compression,
+            wait_for_ready,
+            timeout,
+            metadata,
+            _registered_method=True,
+        )
+
     @staticmethod
     def run_task(
         request,

indexify/functions_sdk/graph.py

@@ -37,6 +37,7 @@ from .indexify_functions import (
     IndexifyRouter,
     RouterCallResult,
 )
+from .invocation_state.local_invocation_state import LocalInvocationState
 from .object_serializer import get_serializer
 
 RouterFn = Annotated[
@@ -236,13 +237,12 @@ class Graph:
                 payload=serializer.serialize(v), encoder=node.input_encoder
             )
         self._results[input.id] = outputs
-        ctx = GraphInvocationContext(
+        self._local_graph_ctx = GraphInvocationContext(
             invocation_id=input.id,
             graph_name=self.name,
             graph_version="1",
-            indexify_client=None,
+            invocation_state=LocalInvocationState(),
         )
-        self._local_graph_ctx = ctx
         self._run(input, outputs)
         return input.id
 

indexify/functions_sdk/image.py

@@ -1,20 +1,88 @@
+import datetime
 import hashlib
 import importlib
+import logging
+import os
+import pathlib
 import sys
+import tarfile
+from io import BytesIO
 from typing import List, Optional
 
+import docker
+import docker.api.build
 from pydantic import BaseModel
 
 
 # Pydantic object for API
 class ImageInformation(BaseModel):
     image_name: str
-    tag: str
-    base_image: str
-    run_strs: List[str]
+    image_hash: str
     image_url: Optional[str] = ""
     sdk_version: str
 
+    # These are deprecated and here for backwards compatibility
+    run_strs: List[str] | None = []
+    tag: str | None = ""
+    base_image: str | None = ""
+
+
+HASH_BUFF_SIZE = 1024**2
+
+
+class BuildOp(BaseModel):
+    op_type: str
+    args: List[str]
+
+    def hash(self, hash):
+        match self.op_type:
+            case "RUN":
+                hash.update("RUN".encode())
+                for a in self.args:
+                    hash.update(a.encode())
+
+            case "COPY":
+                hash.update("COPY".encode())
+                for root, dirs, files in os.walk(self.args[0]):
+                    for file in files:
+                        filename = pathlib.Path(root, file)
+                        with open(filename, "rb") as fp:
+                            data = fp.read(HASH_BUFF_SIZE)
+                            while data:
+                                hash.update(data)
+                                data = fp.read(HASH_BUFF_SIZE)
+
+            case _:
+                raise ValueError(f"Unsupported build op type {self.op_type}")
+
+    def render(self):
+        match self.op_type:
+            case "RUN":
+                return f"RUN {''.join(self.args)}"
+            case "COPY":
+                return f"COPY {self.args[0]} {self.args[1]}"
+            case _:
+                raise ValueError(f"Unsupported build op type {self.op_type}")
+
+
+class Build(BaseModel):
+    """
+    Model for talking with the build service.
+    """
+
+    id: int | None = None
+    namespace: str
+    image_name: str
+    image_hash: str
+    status: str | None
+    result: str | None
+
+    created_at: datetime.datetime | None
+    started_at: datetime.datetime | None = None
+    build_completed_at: datetime.datetime | None = None
+    push_completed_at: datetime.datetime | None = None
+    uri: str | None = None
+
 
 class Image:
     def __init__(self):
@@ -22,7 +90,7 @@ class Image:
         self._tag = "latest"
         self._base_image = BASE_IMAGE_NAME
         self._python_version = LOCAL_PYTHON_VERSION
-        self._run_strs = []
+        self._build_ops = []  # List of ImageOperation
         self._sdk_version = importlib.metadata.version("indexify")
 
     def name(self, image_name):
@@ -38,16 +106,79 @@ class Image:
         return self
 
     def run(self, run_str):
-        self._run_strs.append(run_str)
+        self._build_ops.append(BuildOp(op_type="RUN", args=[run_str]))
+        return self
+
+    def copy(self, source: str, dest: str):
+        self._build_ops.append(BuildOp(op_type="COPY", args=[source, dest]))
         return self
 
     def to_image_information(self):
         return ImageInformation(
             image_name=self._image_name,
-            tag=self._tag,
-            base_image=self._base_image,
-            run_strs=self._run_strs,
             sdk_version=self._sdk_version,
+            image_hash=self.hash(),
+        )
+
+    def build_context(self, filename: str):
+        with tarfile.open(filename, "w:gz") as tf:
+            for op in self._build_ops:
+                if op.op_type == "COPY":
+                    src = op.args[0]
+                    logging.info(f"Adding {src}")
+                    tf.add(src, src)
+
+            dockerfile = self._generate_dockerfile()
+            tarinfo = tarfile.TarInfo("Dockerfile")
+            tarinfo.size = len(dockerfile)
+
+            tf.addfile(tarinfo, BytesIO(dockerfile.encode()))
+
+    def _generate_dockerfile(self, python_sdk_path: Optional[str] = None):
+        docker_contents = [
+            f"FROM {self._base_image}",
+            "RUN mkdir -p ~/.indexify",
+            f"RUN echo {self._image_name} > ~/.indexify/image_name",
+            f"RUN echo {self.hash()} > ~/.indexify/image_hash",
+            "WORKDIR /app",
+        ]
+
+        for build_op in self._build_ops:
+            docker_contents.append(build_op.render())
+
+        if python_sdk_path is not None:
+            logging.info(
+                f"Building image {self._image_name} with local version of the SDK"
+            )
+            if not os.path.exists(python_sdk_path):
+                print(f"error: {python_sdk_path} does not exist")
+                os.exit(1)
+            docker_contents.append(f"COPY {python_sdk_path} /app/python-sdk")
+            docker_contents.append("RUN (cd /app/python-sdk && pip install .)")
+        else:
+            docker_contents.append(f"RUN pip install indexify=={self._sdk_version}")
+
+        docker_file = "\n".join(docker_contents)
+        return docker_file
+
+    def build(self, python_sdk_path: Optional[str] = None, docker_client=None):
+        if docker_client is None:
+            docker_client = docker.from_env()
+            docker_client.ping()
+
+        docker_file = self._generate_dockerfile(python_sdk_path=python_sdk_path)
+        image_name = f"{self._image_name}:{self._tag}"
+
+        docker.api.build.process_dockerfile = lambda dockerfile, path: (
+            "Dockerfile",
+            dockerfile,
+        )
+
+        return docker_client.images.build(
+            path=".",
+            dockerfile=docker_file,
+            tag=image_name,
+            rm=True,
         )
 
     def hash(self) -> str:
@@ -55,7 +186,9 @@ class Image:
             self._image_name.encode()
         )  # Make a hash of the image name
         hash.update(self._base_image.encode())
-        hash.update("".join(self._run_strs).encode())
+        for op in self._build_ops:
+            op.hash(hash)
+
         hash.update(self._sdk_version.encode())
 
         return hash.hexdigest()

indexify/functions_sdk/indexify_functions.py

@@ -1,8 +1,8 @@
 import inspect
 import traceback
+from inspect import Parameter
 from typing import (
     Any,
-    Callable,
     Dict,
     List,
     Optional,
@@ -13,34 +13,27 @@ from typing import (
     get_origin,
 )
 
-from pydantic import BaseModel, Field, PrivateAttr
+from pydantic import BaseModel
+from typing_extensions import get_type_hints
 
 from .data_objects import IndexifyData
 from .image import DEFAULT_IMAGE, Image
+from .invocation_state.invocation_state import InvocationState
 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
-        )
+class GraphInvocationContext:
+    def __init__(
+        self,
+        invocation_id: str,
+        graph_name: str,
+        graph_version: str,
+        invocation_state: InvocationState,
+    ):
+        self.invocation_id = invocation_id
+        self.graph_name = graph_name
+        self.graph_version = graph_version
+        self.invocation_state = invocation_state
 
 
 def is_pydantic_model_from_annotation(type_annotation):
@@ -89,10 +82,22 @@ class IndexifyFunction:
     def run(self, *args, **kwargs) -> Union[List[Any], Any]:
         pass
 
-    def partial(self, **kwargs) -> Callable:
-        from functools import partial
-
-        return partial(self.run, **kwargs)
+    def _call_run(self, *args, **kwargs) -> Union[List[Any], Any]:
+        # Process dictionary argument mapping it to args or to kwargs.
+        if self.accumulate and len(args) == 2 and isinstance(args[1], dict):
+            sig = inspect.signature(self.run)
+            new_args = [args[0]]  # Keep the accumulate argument
+            dict_arg = args[1]
+            new_args_from_dict, new_kwargs = _process_dict_arg(dict_arg, sig)
+            new_args.extend(new_args_from_dict)
+            return self.run(*new_args, **new_kwargs)
+        elif len(args) == 1 and isinstance(args[0], dict):
+            sig = inspect.signature(self.run)
+            dict_arg = args[0]
+            new_args, new_kwargs = _process_dict_arg(dict_arg, sig)
+            return self.run(*new_args, **new_kwargs)
+
+        return self.run(*args, **kwargs)
 
     @classmethod
     def deserialize_output(cls, output: IndexifyData) -> Any:
@@ -111,10 +116,16 @@ class IndexifyRouter:
     def run(self, *args, **kwargs) -> Optional[List[IndexifyFunction]]:
         pass
 
+    # Create run method that preserves signature
+    def _call_run(self, *args, **kwargs):
+        # Process dictionary argument mapping it to args or to kwargs.
+        if len(args) == 1 and isinstance(args[0], dict):
+            sig = inspect.signature(self.run)
+            dict_arg = args[0]
+            new_args, new_kwargs = _process_dict_arg(dict_arg, sig)
+            return self.run(*new_args, **new_kwargs)
 
-from inspect import Parameter, signature
-
-from typing_extensions import get_type_hints
+        return self.run(*args, **kwargs)
 
 
 def _process_dict_arg(dict_arg: dict, sig: inspect.Signature) -> Tuple[list, dict]:
@@ -147,25 +158,6 @@ def indexify_router(
     output_encoder: Optional[str] = "cloudpickle",
 ):
     def construct(fn):
-        # Get function signature using inspect.signature
-        fn_sig = signature(fn)
-        fn_hints = get_type_hints(fn)
-
-        # Create run method that preserves signature
-        def run(self, *args, **kwargs):
-            # Process dictionary argument mapping it to args or to kwargs.
-            if len(args) == 1 and isinstance(args[0], dict):
-                sig = inspect.signature(fn)
-                dict_arg = args[0]
-                new_args, new_kwargs = _process_dict_arg(dict_arg, sig)
-                return fn(*new_args, **new_kwargs)
-
-            return fn(*args, **kwargs)
-
-        # Apply original signature and annotations to run method
-        run.__signature__ = fn_sig
-        run.__annotations__ = fn_hints
-
         attrs = {
             "name": name if name else fn.__name__,
             "description": (
@@ -177,7 +169,7 @@ def indexify_router(
             "placement_constraints": placement_constraints,
             "input_encoder": input_encoder,
             "output_encoder": output_encoder,
-            "run": run,
+            "run": staticmethod(fn),
         }
 
         return type("IndexifyRouter", (IndexifyRouter,), attrs)
@@ -195,32 +187,6 @@ def indexify_function(
     placement_constraints: List[PlacementConstraints] = [],
 ):
     def construct(fn):
-        # Get function signature using inspect.signature
-        fn_sig = signature(fn)
-        fn_hints = get_type_hints(fn)
-
-        # Create run method that preserves signature
-        def run(self, *args, **kwargs):
-            # Process dictionary argument mapping it to args or to kwargs.
-            if self.accumulate and len(args) == 2 and isinstance(args[1], dict):
-                sig = inspect.signature(fn)
-                new_args = [args[0]]  # Keep the accumulate argument
-                dict_arg = args[1]
-                new_args_from_dict, new_kwargs = _process_dict_arg(dict_arg, sig)
-                new_args.extend(new_args_from_dict)
-                return fn(*new_args, **new_kwargs)
-            elif len(args) == 1 and isinstance(args[0], dict):
-                sig = inspect.signature(fn)
-                dict_arg = args[0]
-                new_args, new_kwargs = _process_dict_arg(dict_arg, sig)
-                return fn(*new_args, **new_kwargs)
-
-            return fn(*args, **kwargs)
-
-        # Apply original signature and annotations to run method
-        run.__signature__ = fn_sig
-        run.__annotations__ = fn_hints
-
         attrs = {
             "name": name if name else fn.__name__,
             "description": (
@@ -233,7 +199,7 @@ def indexify_function(
             "accumulate": accumulate,
             "input_encoder": input_encoder,
             "output_encoder": output_encoder,
-            "run": run,
+            "run": staticmethod(fn),
         }
 
         return type("IndexifyFunction", (IndexifyFunction,), attrs)
@@ -303,7 +269,7 @@ class IndexifyFunctionWrapper:
                 args += input
             else:
                 args.append(input)
-            extracted_data = self.indexify_function.run(*args, **kwargs)
+            extracted_data = self.indexify_function._call_run(*args, **kwargs)
         except Exception as e:
             return [], traceback.format_exc()
         if not isinstance(extracted_data, list) and extracted_data is not None:
@@ -330,7 +296,7 @@ class IndexifyFunctionWrapper:
             args.append(input)
 
         try:
-            extracted_data = self.indexify_function.run(*args, **kwargs)
+            extracted_data = self.indexify_function._call_run(*args, **kwargs)
         except Exception as e:
             return [], traceback.format_exc()
         if extracted_data is None:

indexify/functions_sdk/invocation_state/invocation_state.py

@@ -0,0 +1,22 @@
+from typing import Any, Optional
+
+
+class InvocationState:
+    """Abstract interface for Graph invocation state key-value API.
+
+    The API allows to set and get key-value pairs from Indexify functions.
+    The key-value pairs are scoped per Graph invocation.
+    Each new invocation starts with an empty state (empty set of key-value pairs).
+    A value can be any CloudPickleSerializer serializable object."""
+
+    def set(self, key: str, value: Any) -> None:
+        """Set a key-value pair.
+
+        Raises Exception if an error occured."""
+        raise NotImplementedError()
+
+    def get(self, key: str, default: Optional[Any] = None) -> Optional[Any]:
+        """Get a value by key. If the key does not exist, return the default value.
+
+        Raises Exception if an error occured."""
+        raise NotImplementedError()

indexify/functions_sdk/invocation_state/local_invocation_state.py

@@ -0,0 +1,30 @@
+from typing import Any, Dict, Optional
+
+from ..object_serializer import CloudPickleSerializer
+from .invocation_state import InvocationState
+
+
+class LocalInvocationState(InvocationState):
+    """InvocationState that stores the key-value pairs in memory.
+
+    This is intended to be used with local graphs."""
+
+    def __init__(self):
+        """Creates a new instance.
+
+        Caller needs to ensure that the returned instance is only used for a single invocation state.
+        """
+        self._state: Dict[str, bytes] = {}
+
+    def set(self, key: str, value: Any) -> None:
+        # It's important to serialize the value even in the local implementation
+        # so there are no unexpected errors when running in remote graph mode.
+        self._state[key] = CloudPickleSerializer.serialize(value)
+
+    def get(self, key: str, default: Optional[Any] = None) -> Optional[Any]:
+        serialized_value: Optional[bytes] = self._state.get(key, None)
+        return (
+            default
+            if serialized_value is None
+            else CloudPickleSerializer.deserialize(serialized_value)
+        )

indexify/http_client.py

@@ -208,23 +208,6 @@ 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"

{indexify-0.2.44.dist-info → indexify-0.2.45.dist-info}/METADATA

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