indexify 0.3.19__py3-none-any.whl → 0.3.21__py3-none-any.whl

indexify/cli/cli.py

@@ -25,6 +25,9 @@ from rich.theme import Theme
 from tensorlake.functions_sdk.image import Image
 
 from indexify.executor.api_objects import FunctionURI
+from indexify.executor.blob_store.blob_store import BLOBStore
+from indexify.executor.blob_store.local_fs_blob_store import LocalFSBLOBStore
+from indexify.executor.blob_store.s3_blob_store import S3BLOBStore
 from indexify.executor.executor import Executor
 from indexify.executor.executor_flavor import ExecutorFlavor
 from indexify.executor.function_executor.server.subprocess_function_executor_server_factory import (
@@ -197,6 +200,14 @@ def executor(
         )
         exit(1)
 
+    # Enable all available blob stores in OSS because we don't know which one is going to be used.
+    blob_store: BLOBStore = BLOBStore(
+        # Local FS mode is used in tests and in cases when user wants to store data on NFS.
+        local=LocalFSBLOBStore(),
+        # S3 is initiliazed lazily so it's okay to create it even if the user is not going to use it.
+        s3=S3BLOBStore(),
+    )
+
     prometheus_client.Info("cli", "CLI information").info(
         {
             "package": "indexify",
@@ -222,6 +233,7 @@ def executor(
         monitoring_server_host=monitoring_server_host,
         monitoring_server_port=monitoring_server_port,
         enable_grpc_state_reconciler=enable_grpc_state_reconciler,
+        blob_store=blob_store,
     ).run()
 
 

indexify/executor/api_objects.py

@@ -3,6 +3,13 @@ from typing import Any, Dict, List, Optional
 from pydantic import BaseModel
 
 
+class DataPayload(BaseModel):
+    path: str
+    size: int
+    sha256_hash: str
+    content_type: Optional[str] = None
+
+
 class Task(BaseModel):
     id: str
     namespace: str
@@ -16,6 +23,10 @@ class Task(BaseModel):
     "image_uri defines the URI of the image of this task. Optional since some executors do not require it."
     secret_names: Optional[List[str]] = None
     "secret_names defines the names of the secrets to set on function executor. Optional for backward compatibility."
+    graph_payload: Optional[DataPayload] = None
+    input_payload: Optional[DataPayload] = None
+    reducer_input_payload: Optional[DataPayload] = None
+    output_payload_uri_prefix: Optional[str] = None
 
 
 class FunctionURI(BaseModel):
@@ -49,12 +60,6 @@ class TaskResult(BaseModel):
     reducer: bool = False
 
 
-class DataPayload(BaseModel):
-    path: str
-    size: int
-    sha256_hash: str
-
-
 class IngestFnOutputsResponse(BaseModel):
     data_payloads: List[DataPayload]
     stdout: Optional[DataPayload] = None

indexify/executor/blob_store/blob_store.py

@@ -0,0 +1,69 @@
+from typing import Any, Optional
+
+from .local_fs_blob_store import LocalFSBLOBStore
+from .metrics.blob_store import (
+    metric_get_blob_errors,
+    metric_get_blob_latency,
+    metric_get_blob_requests,
+    metric_put_blob_errors,
+    metric_put_blob_latency,
+    metric_put_blob_requests,
+)
+from .s3_blob_store import S3BLOBStore
+
+
+class BLOBStore:
+    """Dispatches generic BLOB store calls to their real backends."""
+
+    def __init__(
+        self, local: Optional[LocalFSBLOBStore] = None, s3: Optional[S3BLOBStore] = None
+    ):
+        """Creates a BLOB store that uses the supplied BLOB stores."""
+        self._local: Optional[LocalFSBLOBStore] = local
+        self._s3: Optional[S3BLOBStore] = s3
+
+    async def get(self, uri: str, logger: Any) -> bytes:
+        """Returns binary value stored in BLOB with the supplied URI.
+
+        Raises Exception on error. Raises KeyError if the BLOB doesn't exist.
+        """
+        with (
+            metric_get_blob_errors.count_exceptions(),
+            metric_get_blob_latency.time(),
+        ):
+            metric_get_blob_requests.inc()
+            if _is_file_uri(uri):
+                self._check_local_is_available()
+                return await self._local.get(uri, logger)
+            else:
+                self._check_s3_is_available()
+                return await self._s3.get(uri, logger)
+
+    async def put(self, uri: str, value: bytes, logger: Any) -> None:
+        """Stores the supplied binary value in a BLOB with the supplied URI.
+
+        Overwrites existing BLOB. Raises Exception on error.
+        """
+        with (
+            metric_put_blob_errors.count_exceptions(),
+            metric_put_blob_latency.time(),
+        ):
+            metric_put_blob_requests.inc()
+            if _is_file_uri(uri):
+                self._check_local_is_available()
+                await self._local.put(uri, value, logger)
+            else:
+                self._check_s3_is_available()
+                await self._s3.put(uri, value, logger)
+
+    def _check_local_is_available(self):
+        if self._local is None:
+            raise RuntimeError("Local file system BLOB store is not available")
+
+    def _check_s3_is_available(self):
+        if self._s3 is None:
+            raise RuntimeError("S3 BLOB store is not available")
+
+
+def _is_file_uri(uri: str) -> bool:
+    return uri.startswith("file://")

indexify/executor/blob_store/local_fs_blob_store.py

@@ -0,0 +1,48 @@
+import asyncio
+import os
+import os.path
+from typing import Any
+
+
+class LocalFSBLOBStore:
+    """BLOB store that stores BLOBs in local file system."""
+
+    async def get(self, uri: str, logger: Any) -> bytes:
+        """Returns binary value stored in file at the supplied URI.
+
+        The URI must be a file URI (starts with "file://"). The path must be absolute.
+        Raises Exception on error. Raises KeyError if the file doesn't exist.
+        """
+        # Run synchronous code in a thread to not block the event loop.
+        return await asyncio.to_thread(self._sync_get, _path_from_file_uri(uri))
+
+    async def put(self, uri: str, value: bytes, logger: Any) -> None:
+        """Stores the supplied binary value in a file at the supplied URI.
+
+        The URI must be a file URI (starts with "file://"). The path must be absolute.
+        Overwrites existing file. Raises Exception on error.
+        """
+        # Run synchronous code in a thread to not block the event loop.
+        return await asyncio.to_thread(self._sync_put, _path_from_file_uri(uri), value)
+
+    def _sync_get(self, path: str) -> bytes:
+        if not os.path.isabs(path):
+            raise ValueError(f"Path {path} must be absolute")
+
+        if os.path.exists(path):
+            with open(path, mode="rb") as blob_file:
+                return blob_file.read()
+        else:
+            raise KeyError(f"File at {path} does not exist")
+
+    def _sync_put(self, path: str, value: bytes) -> None:
+        if not os.path.isabs(path):
+            raise ValueError(f"Path {path} must be absolute")
+
+        os.makedirs(os.path.dirname(path), exist_ok=True)
+        with open(path, mode="wb") as blob_file:
+            blob_file.write(value)
+
+
+def _path_from_file_uri(uri: str) -> str:
+    return uri[7:]  # strip "file://" prefix

indexify/executor/blob_store/metrics/blob_store.py

@@ -0,0 +1,33 @@
+import prometheus_client
+
+from ...monitoring.metrics import latency_metric_for_fast_operation
+
+metric_get_blob_requests: prometheus_client.Counter = prometheus_client.Counter(
+    "get_blob_requests",
+    "Number of get blob requests",
+)
+metric_get_blob_errors: prometheus_client.Counter = prometheus_client.Counter(
+    "get_blob_request_errors",
+    "Number of get blob request errors",
+)
+metric_get_blob_latency: prometheus_client.Histogram = (
+    latency_metric_for_fast_operation(
+        "get_blob_request",
+        "get blob request",
+    )
+)
+
+metric_put_blob_requests: prometheus_client.Counter = prometheus_client.Counter(
+    "put_blob_requests",
+    "Number of put blob requests",
+)
+metric_put_blob_errors: prometheus_client.Counter = prometheus_client.Counter(
+    "put_blob_request_errors",
+    "Number of put blob request errors",
+)
+metric_put_blob_latency: prometheus_client.Histogram = (
+    latency_metric_for_fast_operation(
+        "put_blob_request",
+        "put blob request",
+    )
+)

indexify/executor/blob_store/s3_blob_store.py

@@ -0,0 +1,88 @@
+import asyncio
+from typing import Any, Optional
+
+import boto3
+from botocore.config import Config as BotoConfig
+from botocore.exceptions import ClientError as BotoClientError
+
+_MAX_RETRIES = 3
+
+
+class S3BLOBStore:
+    def __init__(self):
+        self._s3_client: Optional[Any] = None
+
+    def _lazy_create_client(self):
+        """Creates S3 client if it doesn't exist.
+
+        We create the client lazily only if S3 is used.
+        This is because S3 BLOB store is always created by Executor
+        and the creation will fail if user didn't configure S3 credentials and etc.
+        """
+        if self._s3_client is not None:
+            return
+
+        # The credentials and etc are fetched by boto3 library automatically following
+        # https://boto3.amazonaws.com/v1/documentation/api/latest/guide/credentials.html#configuring-credentials
+        # This provides a lot of flexibility for the user and follows a well-known and documented logic.
+        self._s3_client = boto3.client(
+            "s3",
+            config=BotoConfig(
+                # https://boto3.amazonaws.com/v1/documentation/api/latest/guide/retries.html#standard-retry-mode
+                retries={
+                    "max_attempts": _MAX_RETRIES,
+                    "mode": "standard",
+                }
+            ),
+        )
+
+    async def get(self, uri: str, logger: Any) -> bytes:
+        """Returns binary value stored in S3 object at the supplied URI.
+
+        The URI must be S3 URI (starts with "s3://").
+        Raises Exception on error. Raises KeyError if the object doesn't exist.
+        """
+        try:
+            self._lazy_create_client()
+            bucket_name, key = _bucket_name_and_object_key_from_uri(uri)
+            response = await asyncio.to_thread(
+                self._s3_client.get_object, Bucket=bucket_name, Key=key
+            )
+            return response["Body"].read()
+        except BotoClientError as e:
+            logger.error("failed to get S3 object", uri=uri, exc_info=e)
+
+            if e.response["Error"]["Code"] == "NoSuchKey":
+                raise KeyError(f"Object {key} does not exist in bucket {bucket_name}")
+            raise
+        except Exception as e:
+            logger.error("failed to get S3 object", uri=uri, exc_info=e)
+            raise
+
+    async def put(self, uri: str, value: bytes, logger: Any) -> None:
+        """Stores the supplied binary value in a S3 object at the supplied URI.
+
+        The URI must be S3 URI (starts with "s3://").
+        Overwrites existing object. Raises Exception on error.
+        """
+        try:
+            self._lazy_create_client()
+            bucket_name, key = _bucket_name_and_object_key_from_uri(uri)
+            await asyncio.to_thread(
+                self._s3_client.put_object, Bucket=bucket_name, Key=key, Body=value
+            )
+        except Exception as e:
+            logger.error("failed to set S3 object", uri=uri, exc_info=e)
+            raise
+
+
+def _bucket_name_and_object_key_from_uri(uri: str) -> tuple[str, str]:
+    # Example S3 object URI:
+    # s3://test-indexify-server-blob-store-eugene-20250411/225b83f4-2aed-40a7-adee-b7a681f817f2
+    if not uri.startswith("s3://"):
+        raise ValueError(f"S3 URI '{uri}' is missing 's3://' prefix")
+
+    parts = uri[5:].split("/", 1)
+    if len(parts) != 2:
+        raise ValueError(f"Failed parsing bucket name from S3 URI '{uri}'")
+    return parts[0], parts[1]  # bucket_name, key

indexify/executor/downloader.py

@@ -1,13 +1,20 @@
 import asyncio
 import os
-from typing import Any, Optional
+from typing import Any, Optional, Union
 
 import httpx
 import nanoid
 from tensorlake.function_executor.proto.function_executor_pb2 import SerializedObject
+from tensorlake.function_executor.proto.message_validator import MessageValidator
 from tensorlake.utils.http_client import get_httpx_client
 
-from .api_objects import Task
+from indexify.proto.executor_api_pb2 import DataPayload as DataPayloadProto
+from indexify.proto.executor_api_pb2 import (
+    DataPayloadEncoding,
+)
+
+from .api_objects import DataPayload
+from .blob_store.blob_store import BLOBStore
 from .metrics.downloader import (
     metric_graph_download_errors,
     metric_graph_download_latency,
@@ -27,14 +34,24 @@ from .metrics.downloader import (
 
 class Downloader:
     def __init__(
-        self, code_path: str, base_url: str, config_path: Optional[str] = None
+        self,
+        code_path: str,
+        base_url: str,
+        blob_store: BLOBStore,
+        config_path: Optional[str] = None,
     ):
-        self.code_path = code_path
+        self._code_path = code_path
         self._base_url = base_url
         self._client = get_httpx_client(config_path, make_async=True)
+        self._blob_store: BLOBStore = blob_store
 
     async def download_graph(
-        self, namespace: str, graph_name: str, graph_version: str, logger: Any
+        self,
+        namespace: str,
+        graph_name: str,
+        graph_version: str,
+        data_payload: Optional[Union[DataPayload, DataPayloadProto]],
+        logger: Any,
     ) -> SerializedObject:
         logger = logger.bind(module=__name__)
         with (
@@ -47,6 +64,7 @@ class Downloader:
                 namespace=namespace,
                 graph_name=graph_name,
                 graph_version=graph_version,
+                data_payload=data_payload,
                 logger=logger,
             )
 
@@ -56,6 +74,7 @@ class Downloader:
         graph_name: str,
         graph_invocation_id: str,
         input_key: str,
+        data_payload: Optional[DataPayload],
         logger: Any,
     ) -> SerializedObject:
         logger = logger.bind(module=__name__)
@@ -70,6 +89,7 @@ class Downloader:
                 graph_name=graph_name,
                 graph_invocation_id=graph_invocation_id,
                 input_key=input_key,
+                data_payload=data_payload,
                 logger=logger,
             )
 
@@ -80,6 +100,7 @@ class Downloader:
         function_name: str,
         graph_invocation_id: str,
         reducer_output_key: str,
+        data_payload: Optional[Union[DataPayload, DataPayloadProto]],
         logger: Any,
     ) -> SerializedObject:
         logger = logger.bind(module=__name__)
@@ -89,21 +110,27 @@ class Downloader:
             metric_reducer_init_value_download_latency.time(),
         ):
             metric_reducer_init_value_downloads.inc()
-            return await self._fetch_function_init_value(
+            return await self._download_init_value(
                 namespace=namespace,
                 graph_name=graph_name,
                 function_name=function_name,
                 graph_invocation_id=graph_invocation_id,
                 reducer_output_key=reducer_output_key,
+                data_payload=data_payload,
                 logger=logger,
             )
 
     async def _download_graph(
-        self, namespace: str, graph_name: str, graph_version: str, logger: Any
+        self,
+        namespace: str,
+        graph_name: str,
+        graph_version: str,
+        data_payload: Optional[Union[DataPayload, DataPayloadProto]],
+        logger: Any,
     ) -> SerializedObject:
         # Cache graph to reduce load on the server.
         graph_path = os.path.join(
-            self.code_path,
+            self._code_path,
             "graph_cache",
             namespace,
             graph_name,
@@ -118,17 +145,41 @@ class Downloader:
             metric_graphs_from_cache.inc()
             return graph
 
-        graph: SerializedObject = await self._fetch_graph(
-            namespace=namespace,
-            graph_name=graph_name,
-            graph_version=graph_version,
-            logger=logger,
-        )
+        if data_payload is None:
+            graph: SerializedObject = await self._fetch_graph_from_server(
+                namespace=namespace,
+                graph_name=graph_name,
+                graph_version=graph_version,
+                logger=logger,
+            )
+        elif isinstance(data_payload, DataPayloadProto):
+            (
+                MessageValidator(data_payload)
+                .required_field("uri")
+                .required_field("encoding")
+            )
+            data: bytes = await self._blob_store.get(
+                uri=data_payload.uri, logger=logger
+            )
+            return _serialized_object_from_data_payload_proto(
+                data_payload=data_payload,
+                data=data,
+            )
+        elif isinstance(data_payload, DataPayload):
+            data: bytes = await self._blob_store.get(
+                uri=data_payload.path, logger=logger
+            )
+            return _serialized_object_from_data_payload(
+                data_payload=data_payload,
+                data=data,
+            )
+
         # Filesystem operations are synchronous.
         # Run in a separate thread to not block the main event loop.
         # We don't need to wait for the write completion so we use create_task.
         asyncio.create_task(
-            asyncio.to_thread(self._write_cached_graph, graph_path, graph)
+            asyncio.to_thread(self._write_cached_graph, graph_path, graph),
+            name="graph cache write",
         )
 
         return graph
@@ -145,7 +196,7 @@ class Downloader:
             # Another task already cached the graph.
             return None
 
-        tmp_path = os.path.join(self.code_path, "task_graph_cache", nanoid.generate())
+        tmp_path = os.path.join(self._code_path, "task_graph_cache", nanoid.generate())
         os.makedirs(os.path.dirname(tmp_path), exist_ok=True)
         with open(tmp_path, "wb") as f:
             f.write(graph.SerializeToString())
@@ -162,21 +213,87 @@ class Downloader:
         graph_name: str,
         graph_invocation_id: str,
         input_key: str,
+        data_payload: Optional[Union[DataPayload, DataPayloadProto]],
         logger: Any,
     ) -> SerializedObject:
-        first_function_in_graph = graph_invocation_id == input_key.split("|")[-1]
-        if first_function_in_graph:
-            # The first function in Graph gets its input from graph invocation payload.
-            return await self._fetch_graph_invocation_payload(
+        if data_payload is None:
+            first_function_in_graph = graph_invocation_id == input_key.split("|")[-1]
+            if first_function_in_graph:
+                # The first function in Graph gets its input from graph invocation payload.
+                return await self._fetch_graph_invocation_payload_from_server(
+                    namespace=namespace,
+                    graph_name=graph_name,
+                    graph_invocation_id=graph_invocation_id,
+                    logger=logger,
+                )
+            else:
+                return await self._fetch_function_input_from_server(
+                    input_key=input_key, logger=logger
+                )
+        elif isinstance(data_payload, DataPayloadProto):
+            (
+                MessageValidator(data_payload)
+                .required_field("uri")
+                .required_field("encoding")
+            )
+            data: bytes = await self._blob_store.get(
+                uri=data_payload.uri, logger=logger
+            )
+            return _serialized_object_from_data_payload_proto(
+                data_payload=data_payload,
+                data=data,
+            )
+        elif isinstance(data_payload, DataPayload):
+            data: bytes = await self._blob_store.get(
+                uri=data_payload.path, logger=logger
+            )
+            return _serialized_object_from_data_payload(
+                data_payload=data_payload,
+                data=data,
+            )
+
+    async def _download_init_value(
+        self,
+        namespace: str,
+        graph_name: str,
+        function_name: str,
+        graph_invocation_id: str,
+        reducer_output_key: str,
+        data_payload: Optional[Union[DataPayload, DataPayloadProto]],
+        logger: Any,
+    ) -> SerializedObject:
+        if data_payload is None:
+            return await self._fetch_function_init_value_from_server(
                 namespace=namespace,
                 graph_name=graph_name,
+                function_name=function_name,
                 graph_invocation_id=graph_invocation_id,
+                reducer_output_key=reducer_output_key,
                 logger=logger,
             )
-        else:
-            return await self._fetch_function_input(input_key=input_key, logger=logger)
+        elif isinstance(data_payload, DataPayloadProto):
+            (
+                MessageValidator(data_payload)
+                .required_field("uri")
+                .required_field("encoding")
+            )
+            data: bytes = await self._blob_store.get(
+                uri=data_payload.uri, logger=logger
+            )
+            return _serialized_object_from_data_payload_proto(
+                data_payload=data_payload,
+                data=data,
+            )
+        elif isinstance(data_payload, DataPayload):
+            data: bytes = await self._blob_store.get(
+                uri=data_payload.path, logger=logger
+            )
+            return _serialized_object_from_data_payload(
+                data_payload=data_payload,
+                data=data,
+            )
 
-    async def _fetch_graph(
+    async def _fetch_graph_from_server(
         self, namespace: str, graph_name: str, graph_version: str, logger: Any
     ) -> SerializedObject:
         """Downloads the compute graph for the task and returns it."""
@@ -186,7 +303,7 @@ class Downloader:
             logger=logger,
         )
 
-    async def _fetch_graph_invocation_payload(
+    async def _fetch_graph_invocation_payload_from_server(
         self, namespace: str, graph_name: str, graph_invocation_id: str, logger: Any
     ) -> SerializedObject:
         return await self._fetch_url(
@@ -195,7 +312,7 @@ class Downloader:
             logger=logger,
         )
 
-    async def _fetch_function_input(
+    async def _fetch_function_input_from_server(
         self, input_key: str, logger: Any
     ) -> SerializedObject:
         return await self._fetch_url(
@@ -204,7 +321,7 @@ class Downloader:
             logger=logger,
         )
 
-    async def _fetch_function_init_value(
+    async def _fetch_function_init_value_from_server(
         self,
         namespace: str,
         graph_name: str,
@@ -223,7 +340,11 @@ class Downloader:
     async def _fetch_url(
         self, url: str, resource_description: str, logger: Any
     ) -> SerializedObject:
-        logger.info(f"fetching {resource_description}", url=url)
+        logger.warning(
+            f"downloading resource from Server",
+            url=url,
+            resource_description=resource_description,
+        )
         response: httpx.Response = await self._client.get(url)
         try:
             response.raise_for_status()
@@ -252,3 +373,47 @@ def serialized_object_from_http_response(response: httpx.Response) -> Serialized
         return SerializedObject(
             string=response.text, content_type=response.headers["content-type"]
         )
+
+
+def _serialized_object_from_data_payload(
+    data_payload: DataPayload, data: bytes
+) -> SerializedObject:
+    """Converts the given data payload and its data into SerializedObject accepted by Function Executor."""
+    if data_payload.content_type in [
+        "application/octet-stream",
+        "application/pickle",
+    ]:
+        return SerializedObject(bytes=data, content_type=data_payload.content_type)
+    else:
+        return SerializedObject(
+            string=data.decode("utf-8"), content_type=data_payload.content_type
+        )
+
+
+def _serialized_object_from_data_payload_proto(
+    data_payload: DataPayloadProto, data: bytes
+) -> SerializedObject:
+    """Converts the given data payload and its data into SerializedObject accepted by Function Executor.
+
+    Raises ValueError if the supplied data payload can't be converted into serialized object.
+    """
+    if data_payload.encoding == DataPayloadEncoding.DATA_PAYLOAD_ENCODING_BINARY_PICKLE:
+        return SerializedObject(
+            bytes=data,
+            content_type="application/octet-stream",
+        )
+    elif data_payload.encoding == DataPayloadEncoding.DATA_PAYLOAD_ENCODING_UTF8_TEXT:
+        return SerializedObject(
+            content_type="text/plain",
+            string=data.decode("utf-8"),
+        )
+    elif data_payload.encoding == DataPayloadEncoding.DATA_PAYLOAD_ENCODING_UTF8_JSON:
+        result = SerializedObject(
+            content_type="application/json",
+            string=data.decode("utf-8"),
+        )
+        return result
+
+    raise ValueError(
+        f"Can't convert data payload {data_payload} into serialized object"
+    )