indexify 0.3.31__py3-none-any.whl → 0.4.2__py3-none-any.whl

indexify/executor/downloader.py

@@ -1,417 +0,0 @@
-import asyncio
-import os
-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 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,
-    metric_graph_downloads,
-    metric_graphs_from_cache,
-    metric_reducer_init_value_download_errors,
-    metric_reducer_init_value_download_latency,
-    metric_reducer_init_value_downloads,
-    metric_task_input_download_errors,
-    metric_task_input_download_latency,
-    metric_task_input_downloads,
-    metric_tasks_downloading_graphs,
-    metric_tasks_downloading_inputs,
-    metric_tasks_downloading_reducer_init_value,
-)
-
-
-class Downloader:
-    def __init__(
-        self,
-        code_path: str,
-        base_url: str,
-        blob_store: BLOBStore,
-        config_path: Optional[str] = None,
-    ):
-        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,
-        data_payload: Optional[Union[DataPayload, DataPayloadProto]],
-        logger: Any,
-    ) -> SerializedObject:
-        logger = logger.bind(module=__name__)
-        with (
-            metric_graph_download_errors.count_exceptions(),
-            metric_tasks_downloading_graphs.track_inprogress(),
-            metric_graph_download_latency.time(),
-        ):
-            metric_graph_downloads.inc()
-            return await self._download_graph(
-                namespace=namespace,
-                graph_name=graph_name,
-                graph_version=graph_version,
-                data_payload=data_payload,
-                logger=logger,
-            )
-
-    async def download_input(
-        self,
-        namespace: str,
-        graph_name: str,
-        graph_invocation_id: str,
-        input_key: str,
-        data_payload: Optional[DataPayload],
-        logger: Any,
-    ) -> SerializedObject:
-        logger = logger.bind(module=__name__)
-        with (
-            metric_task_input_download_errors.count_exceptions(),
-            metric_tasks_downloading_inputs.track_inprogress(),
-            metric_task_input_download_latency.time(),
-        ):
-            metric_task_input_downloads.inc()
-            return await self._download_input(
-                namespace=namespace,
-                graph_name=graph_name,
-                graph_invocation_id=graph_invocation_id,
-                input_key=input_key,
-                data_payload=data_payload,
-                logger=logger,
-            )
-
-    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:
-        logger = logger.bind(module=__name__)
-        with (
-            metric_reducer_init_value_download_errors.count_exceptions(),
-            metric_tasks_downloading_reducer_init_value.track_inprogress(),
-            metric_reducer_init_value_download_latency.time(),
-        ):
-            metric_reducer_init_value_downloads.inc()
-            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,
-        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,
-            "graph_cache",
-            namespace,
-            graph_name,
-            graph_version,
-        )
-        # Filesystem operations are synchronous.
-        # Run in a separate thread to not block the main event loop.
-        graph: Optional[SerializedObject] = await asyncio.to_thread(
-            self._read_cached_graph, graph_path
-        )
-        if graph is not None:
-            metric_graphs_from_cache.inc()
-            return graph
-
-        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),
-            name="graph cache write",
-        )
-
-        return graph
-
-    def _read_cached_graph(self, path: str) -> Optional[SerializedObject]:
-        if not os.path.exists(path):
-            return None
-
-        with open(path, "rb") as f:
-            return SerializedObject.FromString(f.read())
-
-    def _write_cached_graph(self, path: str, graph: SerializedObject) -> None:
-        if os.path.exists(path):
-            # Another task already cached the graph.
-            return None
-
-        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())
-        os.makedirs(os.path.dirname(path), exist_ok=True)
-        # Atomically rename the fully written file at tmp path.
-        # This allows us to not use any locking because file link/unlink
-        # are atomic operations at filesystem level.
-        # This also allows to share the same cache between multiple Executors.
-        os.replace(tmp_path, path)
-
-    async def _download_input(
-        self,
-        namespace: str,
-        graph_name: str,
-        graph_invocation_id: str,
-        input_key: str,
-        data_payload: Optional[Union[DataPayload, DataPayloadProto]],
-        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
-                )
-        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,
-            )
-        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_from_server(
-        self, namespace: str, graph_name: str, graph_version: str, logger: Any
-    ) -> SerializedObject:
-        """Downloads the compute graph for the task and returns it."""
-        return await self._fetch_url(
-            url=f"{self._base_url}/internal/namespaces/{namespace}/compute_graphs/{graph_name}/versions/{graph_version}/code",
-            resource_description=f"compute graph: {graph_name}",
-            logger=logger,
-        )
-
-    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(
-            url=f"{self._base_url}/namespaces/{namespace}/compute_graphs/{graph_name}/invocations/{graph_invocation_id}/payload",
-            resource_description=f"graph invocation payload: {graph_invocation_id}",
-            logger=logger,
-        )
-
-    async def _fetch_function_input_from_server(
-        self, input_key: str, logger: Any
-    ) -> SerializedObject:
-        return await self._fetch_url(
-            url=f"{self._base_url}/internal/fn_outputs/{input_key}",
-            resource_description=f"function input: {input_key}",
-            logger=logger,
-        )
-
-    async def _fetch_function_init_value_from_server(
-        self,
-        namespace: str,
-        graph_name: str,
-        function_name: str,
-        graph_invocation_id: str,
-        reducer_output_key: str,
-        logger: Any,
-    ) -> SerializedObject:
-        return await self._fetch_url(
-            url=f"{self._base_url}/namespaces/{namespace}/compute_graphs/{graph_name}"
-            f"/invocations/{graph_invocation_id}/fn/{function_name}/output/{reducer_output_key}",
-            resource_description=f"reducer output: {reducer_output_key}",
-            logger=logger,
-        )
-
-    async def _fetch_url(
-        self, url: str, resource_description: str, logger: Any
-    ) -> SerializedObject:
-        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()
-        except httpx.HTTPStatusError as e:
-            logger.error(
-                f"failed to download {resource_description}",
-                error=response.text,
-                exc_info=e,
-            )
-            raise
-
-        return serialized_object_from_http_response(response)
-
-
-def serialized_object_from_http_response(response: httpx.Response) -> SerializedObject:
-    # We're hardcoding the content type currently used by Python SDK. It might change in the future.
-    # There's no other way for now to determine if the response is a bytes or string.
-    if response.headers["content-type"] in [
-        "application/octet-stream",
-        "application/pickle",
-    ]:
-        return SerializedObject(
-            bytes=response.content, content_type=response.headers["content-type"]
-        )
-    else:
-        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"
-    )

indexify/executor/executor_flavor.py

@@ -1,7 +0,0 @@
-from enum import Enum
-
-
-class ExecutorFlavor(Enum):
-    UNKNOWN = "unknown"
-    OSS = "oss"
-    PLATFORM = "platform"

indexify/executor/function_executor/function_executor_state.py

@@ -1,107 +0,0 @@
-import asyncio
-from typing import Any, List, Optional
-
-from .function_executor import FunctionExecutor
-from .function_executor_status import FunctionExecutorStatus, is_status_change_allowed
-from .metrics.function_executor_state import (
-    metric_function_executor_state_not_locked_errors,
-    metric_function_executors_with_status,
-)
-
-
-class FunctionExecutorState:
-    """State of a Function Executor with a particular ID.
-
-    The Function Executor might not exist, i.e. not yet created or destroyed.
-    This object represents all such states. Any state modification must be done
-    under the lock.
-    """
-
-    def __init__(
-        self,
-        id: str,
-        namespace: str,
-        graph_name: str,
-        graph_version: str,
-        function_name: str,
-        image_uri: Optional[str],
-        secret_names: List[str],
-        logger: Any,
-    ):
-        # Read only fields.
-        self.id: str = id
-        self.namespace: str = namespace
-        self.graph_name: str = graph_name
-        self.function_name: str = function_name
-        self.image_uri: Optional[str] = image_uri
-        self.secret_names: List[str] = secret_names
-        self._logger: Any = logger.bind(
-            module=__name__,
-            function_executor_id=id,
-            namespace=namespace,
-            graph_name=graph_name,
-            graph_version=graph_version,
-            function_name=function_name,
-            image_uri=image_uri,
-        )
-        # The lock must be held while modifying the fields below.
-        self.lock: asyncio.Lock = asyncio.Lock()
-        # TODO: Move graph_version to immutable fields once we migrate to gRPC State Reconciler.
-        self.graph_version: str = graph_version
-        self.status: FunctionExecutorStatus = FunctionExecutorStatus.DESTROYED
-        self.status_change_notifier: asyncio.Condition = asyncio.Condition(
-            lock=self.lock
-        )
-        self.function_executor: Optional[FunctionExecutor] = None
-        metric_function_executors_with_status.labels(status=self.status.name).inc()
-
-    async def wait_status(self, allowlist: List[FunctionExecutorStatus]) -> None:
-        """Waits until Function Executor status reaches one of the allowed values.
-
-        The caller must hold the lock.
-        """
-        self.check_locked()
-        while self.status not in allowlist:
-            await self.status_change_notifier.wait()
-
-    async def set_status(self, new_status: FunctionExecutorStatus) -> None:
-        """Sets the status of the Function Executor.
-
-        The caller must hold the lock.
-        Raises ValueError if the status change is not allowed.
-        """
-        self.check_locked()
-        if is_status_change_allowed(self.status, new_status):
-            # If status didn't change then still log it for visibility.
-            self._logger.info(
-                "function executor status changed",
-                old_status=self.status.name,
-                new_status=new_status.name,
-            )
-            metric_function_executors_with_status.labels(status=self.status.name).dec()
-            metric_function_executors_with_status.labels(status=new_status.name).inc()
-            self.status = new_status
-            self.status_change_notifier.notify_all()
-        else:
-            raise ValueError(
-                f"Invalid status change from {self.status} to {new_status}"
-            )
-
-    # TODO: Delete this method once HTTP protocol is removed as it's used only there.
-    async def destroy_function_executor(self) -> None:
-        """Destroys the Function Executor if it exists.
-
-        The caller must hold the lock.
-        """
-        self.check_locked()
-        await self.set_status(FunctionExecutorStatus.DESTROYING)
-        if self.function_executor is not None:
-            await self.function_executor.destroy()
-            self.function_executor = None
-        await self.set_status(FunctionExecutorStatus.DESTROYED)
-
-    def check_locked(self) -> None:
-        """Raises an exception if the lock is not held."""
-        if not self.lock.locked():
-            metric_function_executor_state_not_locked_errors.inc()
-            raise RuntimeError("The FunctionExecutorState lock must be held.")

indexify/executor/function_executor/function_executor_states_container.py

@@ -1,93 +0,0 @@
-import asyncio
-from typing import Any, AsyncGenerator, Dict, List, Optional
-
-from .function_executor_state import FunctionExecutorState
-from .function_executor_status import FunctionExecutorStatus
-from .metrics.function_executor_state_container import (
-    metric_function_executor_states_count,
-)
-
-
-class FunctionExecutorStatesContainer:
-    """An asyncio concurrent container for the function executor states."""
-
-    def __init__(self, logger: Any):
-        # The fields below are protected by the lock.
-        self._lock: asyncio.Lock = asyncio.Lock()
-        self._states: Dict[str, FunctionExecutorState] = {}
-        self._is_shutdown: bool = False
-        self._logger: Any = logger.bind(module=__name__)
-
-    async def get_or_create_state(
-        self,
-        id: str,
-        namespace: str,
-        graph_name: str,
-        graph_version: str,
-        function_name: str,
-        image_uri: Optional[str],
-        secret_names: List[str],
-    ) -> FunctionExecutorState:
-        """Get or create a function executor state with the given ID.
-
-        If the state already exists, it is returned. Otherwise, a new state is created from the supplied task.
-        Raises Exception if it's not possible to create a new state at this time."""
-        async with self._lock:
-            if self._is_shutdown:
-                raise RuntimeError(
-                    "Function Executor states container is shutting down."
-                )
-
-            if id not in self._states:
-                state = FunctionExecutorState(
-                    id=id,
-                    namespace=namespace,
-                    graph_name=graph_name,
-                    graph_version=graph_version,
-                    function_name=function_name,
-                    image_uri=image_uri,
-                    secret_names=secret_names,
-                    logger=self._logger,
-                )
-                self._states[id] = state
-                metric_function_executor_states_count.set(len(self._states))
-
-            return self._states[id]
-
-    async def get(self, id: str) -> FunctionExecutorState:
-        """Get the state with the given ID. Raises Exception if the state does not exist."""
-        async with self._lock:
-            return self._states[id]
-
-    async def __aiter__(self) -> AsyncGenerator[FunctionExecutorState, None]:
-        async with self._lock:
-            for state in self._states.values():
-                yield state
-
-    async def pop(self, id: str) -> FunctionExecutorState:
-        """Removes the state with the given ID and returns it."""
-        async with self._lock:
-            state = self._states.pop(id)
-            metric_function_executor_states_count.set(len(self._states))
-            return state
-
-    def exists(self, id: str) -> bool:
-        """Check if the state with the given ID exists."""
-        return id in self._states
-
-    async def shutdown(self):
-        # Function Executors are outside the Executor process
-        # so they need to get cleaned up explicitly and reliably.
-        async with self._lock:
-            self._is_shutdown = True  # No new Function Executor States can be created.
-            while self._states:
-                id, state = self._states.popitem()
-                metric_function_executor_states_count.set(len(self._states))
-                # Only ongoing tasks who have a reference to the state already can see it.
-                # The state is unlocked while a task is running inside Function Executor.
-                async with state.lock:
-                    await state.set_status(FunctionExecutorStatus.SHUTDOWN)
-                    if state.function_executor is not None:
-                        await state.function_executor.destroy()
-                        state.function_executor = None
-                        # The task running inside the Function Executor will fail because it's destroyed.