indexify 0.2.48__py3-none-any.whl → 0.3.1__py3-none-any.whl

indexify/functions_sdk/image.py

@@ -1,222 +0,0 @@
-import datetime
-import hashlib
-import importlib
-import logging
-import os
-import pathlib
-import sys
-import tarfile
-from io import BytesIO
-from typing import Dict, List, Optional
-
-import docker
-import docker.api.build
-from pydantic import BaseModel
-
-
-# Pydantic object for API
-class ImageInformation(BaseModel):
-    image_name: 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
-    options: Dict[str, str] = {}
-    args: List[str]
-
-    def hash(self, hash):
-        match self.op_type:
-            case "RUN" | "ADD":
-                hash.update(self.op_type.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" | "ADD":
-                options = [f"--{k}={v}" for k, v in self.options.items()]
-                return f"{self.op_type} {' '.join(options)} {' '.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
-    error_message: str | None = None  # Only provided when result is "failed"
-
-    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):
-        self._image_name = None
-        self._tag = "latest"
-        self._base_image = BASE_IMAGE_NAME
-        self._python_version = LOCAL_PYTHON_VERSION
-        self._build_ops = []  # List of ImageOperation
-        self._sdk_version = importlib.metadata.version("indexify")
-
-    def name(self, image_name):
-        self._image_name = image_name
-        return self
-
-    def tag(self, tag):
-        self._tag = tag
-        return self
-
-    def base_image(self, base_image):
-        self._base_image = base_image
-        return self
-
-    def add(self, source: str, dest: str, **kwargs):
-        self._build_ops.append(
-            BuildOp(op_type="ADD", args=[source, dest], options=kwargs)
-        )
-        return self
-
-    def run(self, run_str, **kwargs):
-        self._build_ops.append(BuildOp(op_type="RUN", args=[run_str], options=kwargs))
-        return self
-
-    def copy(self, source: str, dest: str, **kwargs):
-        self._build_ops.append(
-            BuildOp(op_type="COPY", args=[source, dest], options=kwargs)
-        )
-        return self
-
-    def to_image_information(self):
-        return ImageInformation(
-            image_name=self._image_name,
-            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:
-        hash = hashlib.sha256(
-            self._image_name.encode()
-        )  # Make a hash of the image name
-        hash.update(self._base_image.encode())
-        for op in self._build_ops:
-            op.hash(hash)
-
-        hash.update(self._sdk_version.encode())
-
-        return hash.hexdigest()
-
-
-LOCAL_PYTHON_VERSION = f"{sys.version_info.major}.{sys.version_info.minor}"
-BASE_IMAGE_NAME = f"python:{LOCAL_PYTHON_VERSION}-slim-bookworm"
-
-
-def GetDefaultPythonImage(python_version: str):
-    return (
-        Image()
-        .name("tensorlake/indexify-executor-default")
-        .base_image(f"python:{python_version}-slim-bookworm")
-        .tag(python_version)
-    )
-
-
-DEFAULT_IMAGE = GetDefaultPythonImage(LOCAL_PYTHON_VERSION)

indexify/functions_sdk/indexify_functions.py

@@ -1,354 +0,0 @@
-import inspect
-import traceback
-from inspect import Parameter
-from typing import (
-    Any,
-    Dict,
-    List,
-    Optional,
-    Tuple,
-    Type,
-    Union,
-    get_args,
-    get_origin,
-)
-
-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:
-    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):
-    # If it's a string representation
-    if isinstance(type_annotation, str):
-        # Extract the class name from the string
-        class_name = type_annotation.split("'")[-2].split(".")[-1]
-        # This part is tricky and might require additional context or imports
-        # You might need to import the actual class or module where it's defined
-        # For example:
-        # from indexify.functions_sdk.data_objects import File
-        # return issubclass(eval(class_name), BaseModel)
-        return False  # Default to False if we can't evaluate
-
-    # If it's a Type object
-    origin = get_origin(type_annotation)
-    if origin is not None:
-        # Handle generic types like List[File], Optional[File], etc.
-        args = get_args(type_annotation)
-        if args:
-            return is_pydantic_model_from_annotation(args[0])
-
-    # If it's a direct class reference
-    if isinstance(type_annotation, type):
-        return issubclass(type_annotation, BaseModel)
-
-    return False
-
-
-class PlacementConstraints(BaseModel):
-    min_python_version: Optional[str] = "3.9"
-    max_python_version: Optional[str] = None
-    platform: Optional[str] = None
-    image_name: Optional[str] = None
-
-
-class IndexifyFunction:
-    name: str = ""
-    description: str = ""
-    image: Optional[Image] = DEFAULT_IMAGE
-    placement_constraints: List[PlacementConstraints] = []
-    accumulate: Optional[Type[Any]] = None
-    input_encoder: Optional[str] = "cloudpickle"
-    output_encoder: Optional[str] = "cloudpickle"
-
-    def run(self, *args, **kwargs) -> Union[List[Any], Any]:
-        pass
-
-    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:
-        serializer = get_serializer(cls.output_encoder)
-        return serializer.deserialize(output.payload)
-
-
-class IndexifyRouter:
-    name: str = ""
-    description: str = ""
-    image: Optional[Image] = DEFAULT_IMAGE
-    placement_constraints: List[PlacementConstraints] = []
-    input_encoder: Optional[str] = "cloudpickle"
-    output_encoder: Optional[str] = "cloudpickle"
-
-    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)
-
-        return self.run(*args, **kwargs)
-
-
-def _process_dict_arg(dict_arg: dict, sig: inspect.Signature) -> Tuple[list, dict]:
-    new_args = []
-    new_kwargs = {}
-    remaining_kwargs = dict_arg.copy()
-
-    # Match dictionary keys to function parameters
-    for param_name, param in sig.parameters.items():
-        if param_name in dict_arg:
-            new_args.append(dict_arg[param_name])
-            remaining_kwargs.pop(param_name, None)
-
-    if any(v.kind == Parameter.VAR_KEYWORD for v in sig.parameters.values()):
-        # Combine remaining dict items with additional kwargs
-        new_kwargs.update(remaining_kwargs)
-    elif len(remaining_kwargs) > 0:
-        # If there are remaining kwargs, add them as a single dict argument
-        new_args.append(remaining_kwargs)
-
-    return new_args, new_kwargs
-
-
-def indexify_router(
-    name: Optional[str] = None,
-    description: Optional[str] = "",
-    image: Optional[Image] = DEFAULT_IMAGE,
-    placement_constraints: List[PlacementConstraints] = [],
-    input_encoder: Optional[str] = "cloudpickle",
-    output_encoder: Optional[str] = "cloudpickle",
-):
-    def construct(fn):
-        attrs = {
-            "name": name if name else fn.__name__,
-            "description": (
-                description
-                if description
-                else (fn.__doc__ or "").strip().replace("\n", "")
-            ),
-            "image": image,
-            "placement_constraints": placement_constraints,
-            "input_encoder": input_encoder,
-            "output_encoder": output_encoder,
-            "run": staticmethod(fn),
-        }
-
-        return type("IndexifyRouter", (IndexifyRouter,), attrs)
-
-    return construct
-
-
-def indexify_function(
-    name: Optional[str] = None,
-    description: Optional[str] = "",
-    image: Optional[Image] = DEFAULT_IMAGE,
-    accumulate: Optional[Type[BaseModel]] = None,
-    input_encoder: Optional[str] = "cloudpickle",
-    output_encoder: Optional[str] = "cloudpickle",
-    placement_constraints: List[PlacementConstraints] = [],
-):
-    def construct(fn):
-        attrs = {
-            "name": name if name else fn.__name__,
-            "description": (
-                description
-                if description
-                else (fn.__doc__ or "").strip().replace("\n", "")
-            ),
-            "image": image,
-            "placement_constraints": placement_constraints,
-            "accumulate": accumulate,
-            "input_encoder": input_encoder,
-            "output_encoder": output_encoder,
-            "run": staticmethod(fn),
-        }
-
-        return type("IndexifyFunction", (IndexifyFunction,), attrs)
-
-    return construct
-
-
-class FunctionCallResult(BaseModel):
-    ser_outputs: List[IndexifyData]
-    traceback_msg: Optional[str] = None
-
-
-class RouterCallResult(BaseModel):
-    edges: List[str]
-    traceback_msg: Optional[str] = None
-
-
-class IndexifyFunctionWrapper:
-    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):
-            raise TypeError("Input must be an instance of IndexifyFunction")
-
-        extract_method = self.indexify_function.run
-        type_hints = get_type_hints(extract_method)
-        return_type = type_hints.get("return", Any)
-        if get_origin(return_type) is list:
-            return_type = get_args(return_type)[0]
-        elif get_origin(return_type) is Union:
-            inner_types = get_args(return_type)
-            if len(inner_types) == 2 and type(None) in inner_types:
-                return_type = (
-                    inner_types[0] if inner_types[1] is type(None) else inner_types[1]
-                )
-        return return_type
-
-    def get_input_types(self) -> Dict[str, Any]:
-        if not isinstance(self.indexify_function, IndexifyFunction):
-            raise TypeError("Input must be an instance of IndexifyFunction")
-
-        extract_method = self.indexify_function.run
-        type_hints = get_type_hints(extract_method)
-        return {
-            k: v
-            for k, v in type_hints.items()
-            if k != "return" and not is_pydantic_model_from_annotation(v)
-        }
-
-    def run_router(
-        self, input: Union[Dict, Type[BaseModel]]
-    ) -> Tuple[List[str], Optional[str]]:
-        args = []
-        kwargs = {}
-        try:
-            # tuple and list are considered positional arguments, list is used for compatibility
-            # with json encoding which won't deserialize in tuple.
-            if isinstance(input, tuple) or isinstance(input, list):
-                args += input
-            elif isinstance(input, dict):
-                kwargs.update(input)
-            else:
-                args.append(input)
-            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:
-            return [extracted_data.name], None
-        edges = []
-        for fn in extracted_data or []:
-            edges.append(fn.name)
-        return edges, None
-
-    def run_fn(
-        self, input: Union[Dict, Type[BaseModel], List, Tuple], acc: Type[Any] = None
-    ) -> Tuple[List[Any], Optional[str]]:
-        args = []
-        kwargs = {}
-
-        if acc is not None:
-            args.append(acc)
-
-        # tuple and list are considered positional arguments, list is used for compatibility
-        # with json encoding which won't deserialize in tuple.
-        if isinstance(input, tuple) or isinstance(input, list):
-            args += input
-        elif isinstance(input, dict):
-            kwargs.update(input)
-        else:
-            args.append(input)
-
-        try:
-            extracted_data = self.indexify_function._call_run(*args, **kwargs)
-        except Exception as e:
-            return [], traceback.format_exc()
-        if extracted_data is None:
-            return [], None
-
-        output = (
-            extracted_data if isinstance(extracted_data, list) else [extracted_data]
-        )
-        return output, None
-
-    def invoke_fn_ser(
-        self, name: str, input: IndexifyData, acc: Optional[Any] = None
-    ) -> FunctionCallResult:
-        input = self.deserialize_input(name, input)
-        input_serializer = get_serializer(self.indexify_function.input_encoder)
-        output_serializer = get_serializer(self.indexify_function.output_encoder)
-        if acc is not None:
-            acc = input_serializer.deserialize(acc.payload)
-        if acc is None and self.indexify_function.accumulate is not None:
-            acc = self.indexify_function.accumulate()
-        outputs, err = self.run_fn(input, acc=acc)
-        ser_outputs = [
-            IndexifyData(
-                payload=output_serializer.serialize(output),
-                encoder=self.indexify_function.output_encoder,
-            )
-            for output in outputs
-        ]
-        return FunctionCallResult(ser_outputs=ser_outputs, traceback_msg=err)
-
-    def invoke_router(self, name: str, input: IndexifyData) -> RouterCallResult:
-        input = self.deserialize_input(name, input)
-        edges, err = self.run_router(input)
-        return RouterCallResult(edges=edges, traceback_msg=err)
-
-    def deserialize_input(self, compute_fn: str, indexify_data: IndexifyData) -> Any:
-        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/invocation_state/invocation_state.py

@@ -1,22 +0,0 @@
-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

@@ -1,30 +0,0 @@
-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/functions_sdk/object_serializer.py

@@ -1,68 +0,0 @@
-import json
-from typing import Any, List, Type
-
-import cloudpickle
-
-
-def get_serializer(serializer_type: str) -> Any:
-    if serializer_type == "cloudpickle":
-        return CloudPickleSerializer()
-    elif serializer_type == "json":
-        return JsonSerializer()
-    elif serializer_type == JsonSerializer.content_type:
-        return JsonSerializer()
-    elif serializer_type == CloudPickleSerializer.content_type:
-        return CloudPickleSerializer()
-    raise ValueError(f"Unknown serializer type: {serializer_type}")
-
-
-class JsonSerializer:
-    content_type = "application/json"
-    encoding_type = "json"
-
-    @staticmethod
-    def serialize(data: Any) -> str:
-        try:
-            return json.dumps(data)
-        except Exception as e:
-            raise ValueError(f"failed to serialize data with json: {e}")
-
-    @staticmethod
-    def deserialize(data: str) -> Any:
-        try:
-            if isinstance(data, bytes):
-                data = data.decode("utf-8")
-            return json.loads(data)
-        except Exception as e:
-            raise ValueError(f"failed to deserialize data with json: {e}")
-
-    @staticmethod
-    def serialize_list(data: List[Any]) -> str:
-        return json.dumps(data)
-
-    @staticmethod
-    def deserialize_list(data: str, t: Type) -> List[Any]:
-        if isinstance(data, bytes):
-            data = data.decode("utf-8")
-        return json.loads(data)
-
-
-class CloudPickleSerializer:
-    content_type = "application/octet-stream"
-    encoding_type = "cloudpickle"
-
-    @staticmethod
-    def serialize(data: Any) -> bytes:
-        return cloudpickle.dumps(data)
-
-    @staticmethod
-    def deserialize(data: bytes) -> Any:
-        return cloudpickle.loads(data)
-
-    @staticmethod
-    def serialize_list(data: List[Any]) -> bytes:
-        return cloudpickle.dumps(data)
-
-    @staticmethod
-    def deserialize_list(data: bytes) -> List[Any]:
-        return cloudpickle.loads(data)