indexify 0.3.19__py3-none-any.whl → 0.3.20__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/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} is not 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} is not 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,85 @@
+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]:
+    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

@@ -5,9 +5,15 @@ from typing import Any, Optional
 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,
+    DataPayloadEncoding,
+)
+
+from .blob_store.blob_store import BLOBStore
 from .metrics.downloader import (
     metric_graph_download_errors,
     metric_graph_download_latency,
@@ -27,14 +33,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[DataPayload],
+        logger: Any,
     ) -> SerializedObject:
         logger = logger.bind(module=__name__)
         with (
@@ -47,6 +63,7 @@ class Downloader:
                 namespace=namespace,
                 graph_name=graph_name,
                 graph_version=graph_version,
+                data_payload=data_payload,
                 logger=logger,
             )
 
@@ -56,6 +73,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 +88,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 +99,7 @@ class Downloader:
         function_name: str,
         graph_invocation_id: str,
         reducer_output_key: str,
+        data_payload: Optional[DataPayload],
         logger: Any,
     ) -> SerializedObject:
         logger = logger.bind(module=__name__)
@@ -89,21 +109,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[DataPayload],
+        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 +144,33 @@ 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,
+            )
+        else:
+            (
+                MessageValidator(data_payload)
+                .required_field("uri")
+                .required_field("encoding")
+            )
+            data: bytes = await self._blob_store.get(
+                uri=data_payload.uri, logger=logger
+            )
+            return _data_payload_to_serialized_object(
+                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 +187,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 +204,71 @@ class Downloader:
         graph_name: str,
         graph_invocation_id: str,
         input_key: str,
+        data_payload: Optional[DataPayload],
+        logger: Any,
+    ) -> SerializedObject:
+        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
+                )
+        else:
+            (
+                MessageValidator(data_payload)
+                .required_field("uri")
+                .required_field("encoding")
+            )
+            data: bytes = await self._blob_store.get(
+                uri=data_payload.uri, logger=logger
+            )
+            return _data_payload_to_serialized_object(
+                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[DataPayload],
         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:
+            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)
+            (
+                MessageValidator(data_payload)
+                .required_field("uri")
+                .required_field("encoding")
+            )
+            data: bytes = await self._blob_store.get(
+                uri=data_payload.uri, logger=logger
+            )
+            return _data_payload_to_serialized_object(
+                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 +278,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 +287,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 +296,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,
@@ -252,3 +344,32 @@ def serialized_object_from_http_response(response: httpx.Response) -> Serialized
         return SerializedObject(
             string=response.text, content_type=response.headers["content-type"]
         )
+
+
+def _data_payload_to_serialized_object(
+    data_payload: DataPayload, 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"
+    )

indexify/executor/executor.py

@@ -12,6 +12,7 @@ from tensorlake.utils.logging import suppress as suppress_logging
 from indexify.proto.executor_api_pb2 import ExecutorStatus
 
 from .api_objects import FunctionURI, Task
+from .blob_store.blob_store import BLOBStore
 from .downloader import Downloader
 from .executor_flavor import ExecutorFlavor
 from .function_executor.function_executor_states_container import (
@@ -69,6 +70,7 @@ class Executor:
         monitoring_server_host: str,
         monitoring_server_port: int,
         enable_grpc_state_reconciler: bool,
+        blob_store: BLOBStore,
     ):
         self._logger = structlog.get_logger(module=__name__)
         self._is_shutdown: bool = False
@@ -95,7 +97,10 @@ class Executor:
             self._function_executor_states
         )
         self._downloader = Downloader(
-            code_path=code_path, base_url=self._base_url, config_path=config_path
+            code_path=code_path,
+            base_url=self._base_url,
+            blob_store=blob_store,
+            config_path=config_path,
         )
         self._function_allowlist: Optional[List[FunctionURI]] = function_allowlist
         self._function_executor_server_factory = function_executor_server_factory
@@ -189,12 +194,15 @@ class Executor:
                 signum, self.shutdown, asyncio.get_event_loop()
             )
 
-        asyncio.get_event_loop().create_task(self._monitoring_server.run())
-        if self._state_reporter is not None:
-            self._state_reporter.update_executor_status(
-                ExecutorStatus.EXECUTOR_STATUS_RUNNING
-            )
-            asyncio.get_event_loop().create_task(self._state_reporter.run())
+        asyncio.get_event_loop().create_task(
+            self._monitoring_server.run(), name="monitoring server runner"
+        )
+        self._state_reporter.update_executor_status(
+            ExecutorStatus.EXECUTOR_STATUS_RUNNING
+        )
+        asyncio.get_event_loop().create_task(
+            self._state_reporter.run(), name="state reporter runner"
+        )
 
         metric_executor_state.state("running")
         self._startup_probe_handler.set_ready()
@@ -215,7 +223,6 @@ class Executor:
         """Runs the gRPC state reconciler and state reporter.
 
         Never raises any exceptions."""
-        asyncio.create_task(self._state_reporter.run())
         await self._state_reconciler.run()
 
     async def _http_task_runner_loop(self):
@@ -224,11 +231,15 @@ class Executor:
                 async for task in self._task_fetcher.run():
                     metric_tasks_fetched.inc()
                     if not self._is_shutdown:
-                        asyncio.create_task(self._run_task(task))
+                        asyncio.create_task(
+                            self._run_task(task), name="task runner (http mode)"
+                        )
+                self._logger.info("fetching tasks finished, reconnecting in 5 seconds")
             except Exception as e:
                 self._logger.error(
                     "failed fetching tasks, retrying in 5 seconds", exc_info=e
                 )
+            if not self._is_shutdown:
                 await asyncio.sleep(5)
 
     async def _run_task(self, task: Task) -> None:
@@ -293,12 +304,14 @@ class Executor:
             graph_name=task.compute_graph,
             graph_version=task.graph_version,
             logger=logger,
+            data_payload=None,
         )
         input: SerializedObject = await self._downloader.download_input(
             namespace=task.namespace,
             graph_name=task.compute_graph,
             graph_invocation_id=task.invocation_id,
             input_key=task.input_key,
+            data_payload=None,
             logger=logger,
         )
         init_value: Optional[SerializedObject] = (
@@ -311,6 +324,7 @@ class Executor:
                     function_name=task.compute_fn,
                     graph_invocation_id=task.invocation_id,
                     reducer_output_key=task.reducer_output_id,
+                    data_payload=None,
                     logger=logger,
                 )
             )
@@ -380,12 +394,12 @@ class Executor:
         if self._task_runner is not None:
             await self._task_runner.shutdown()
 
-        if self._channel_manager is not None:
-            await self._channel_manager.shutdown()
         if self._state_reporter is not None:
             await self._state_reporter.shutdown()
         if self._state_reconciler is not None:
             await self._state_reconciler.shutdown()
+        if self._channel_manager is not None:
+            await self._channel_manager.destroy()
 
         # We need to shutdown all users of FE states first,
         # otherwise states might disappear unexpectedly and we might
@@ -397,7 +411,7 @@ class Executor:
         # The current task is cancelled, the code after this line will not run.
 
     def shutdown(self, loop):
-        loop.create_task(self._shutdown(loop))
+        loop.create_task(self._shutdown(loop), name="executor shutdown")
 
     def _task_logger(self, task: Task) -> Any:
         return self._logger.bind(