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

indexify/executor/runtime_probes.py

@@ -1,68 +0,0 @@
-import os
-import platform
-import sys
-from typing import Any, Dict, Tuple
-
-from pydantic import BaseModel
-
-DEFAULT_EXECUTOR = "tensorlake/indexify-executor-default"
-# Empty string is used as a default hash which tells the scheduler to accept any hash.
-DEFAULT_HASH = ""
-
-
-class ProbeInfo(BaseModel):
-    image_name: str
-    image_hash: str
-    python_major_version: int
-    labels: Dict[str, Any] = {}
-    is_default_executor: bool
-
-
-class RuntimeProbes:
-    def __init__(self) -> None:
-        self._image_name = self._read_image_name()
-        self._image_hash = self._read_image_hash()
-        self._os_name = platform.system()
-        self._architecture = platform.machine()
-        (
-            self._python_version_major,
-            self._python_version_minor,
-        ) = self._get_python_version()
-
-    def _read_image_name(self) -> str:
-        file_path = os.path.expanduser("~/.indexify/image_name")
-        if os.path.exists(file_path):
-            with open(file_path, "r") as file:
-                return file.read().strip()
-        return DEFAULT_EXECUTOR
-
-    def _read_image_hash(self) -> str:
-        file_path = os.path.expanduser("~/.indexify/image_hash")
-        if os.path.exists(file_path):
-            with open(file_path, "r") as file:
-                return file.read().strip()
-        return DEFAULT_HASH
-
-    def _get_python_version(self) -> Tuple[int, int]:
-        version_info = sys.version_info
-        return version_info.major, version_info.minor
-
-    def _is_default_executor(self):
-        return True if self._read_image_name() == DEFAULT_EXECUTOR else False
-
-    def probe(self) -> ProbeInfo:
-        labels = {
-            "os": self._os_name,
-            "image_name": self._image_name,
-            "architecture": self._architecture,
-            "python_major_version": self._python_version_major,
-            "python_minor_version": self._python_version_minor,
-        }
-
-        return ProbeInfo(
-            image_name=self._image_name,
-            image_hash=self._image_hash,
-            python_major_version=self._python_version_major,
-            labels=labels,
-            is_default_executor=self._is_default_executor(),
-        )

indexify/executor/task_fetcher.py

@@ -1,96 +0,0 @@
-import json
-import time
-from socket import gethostname
-from typing import AsyncGenerator, Dict, List, Optional
-
-import structlog
-from httpx_sse import aconnect_sse
-from tensorlake.utils.http_client import get_httpx_client
-
-from .api_objects import ExecutorMetadata, FunctionURI, Task
-from .metrics.task_fetcher import (
-    metric_server_registration_errors,
-    metric_server_registration_latency,
-    metric_server_registrations,
-)
-from .runtime_probes import ProbeInfo, RuntimeProbes
-
-
-class TaskFetcher:
-    """Registers with Indexify server and fetches tasks from it."""
-
-    def __init__(
-        self,
-        executor_id: str,
-        executor_version: str,
-        labels: Dict[str, str],
-        function_allowlist: Optional[List[FunctionURI]],
-        protocol: str,
-        indexify_server_addr: str,
-        config_path: Optional[str] = None,
-    ):
-        self._protocol: str = protocol
-        self._indexify_server_addr: str = indexify_server_addr
-        self.config_path = config_path
-        self._logger = structlog.get_logger(module=__name__)
-
-        probe_info: ProbeInfo = RuntimeProbes().probe()
-        all_labels = probe_info.labels.copy()
-        all_labels.update(labels)
-
-        self._executor_metadata: ExecutorMetadata = ExecutorMetadata(
-            id=executor_id,
-            executor_version=executor_version,
-            addr=gethostname(),
-            function_allowlist=function_allowlist,
-            labels=all_labels,
-        )
-
-    async def run(self) -> AsyncGenerator[Task, None]:
-        """Fetches tasks that Indexify server assigned to the Executor.
-
-        Raises an exception if error occurred."""
-        url = f"{self._protocol}://{self._indexify_server_addr}/internal/executors/{self._executor_metadata.id}/tasks"
-
-        self._logger.info(
-            "registering_executor",
-            executor_id=self._executor_metadata.id,
-            url=url,
-            executor_version=self._executor_metadata.executor_version,
-        )
-        metric_server_registrations.inc()
-        registration_start_time: float = time.monotonic()
-
-        async with get_httpx_client(
-            config_path=self.config_path, make_async=True
-        ) as client:
-            async with aconnect_sse(
-                client,
-                "POST",
-                url,
-                json=self._executor_metadata.model_dump(),
-                headers={"Content-Type": "application/json"},
-            ) as event_source:
-                try:
-                    event_source.response.raise_for_status()
-                except Exception as e:
-                    metric_server_registration_errors.inc()
-                    await event_source.response.aread()
-                    raise Exception(
-                        "failed to register at server. "
-                        f"Response code: {event_source.response.status_code}. "
-                        f"Response text: '{event_source.response.text}'."
-                    ) from e
-                finally:
-                    metric_server_registration_latency.observe(
-                        time.monotonic() - registration_start_time
-                    )
-
-                self._logger.info(
-                    "executor_registered", executor_id=self._executor_metadata.id
-                )
-
-                async for sse in event_source.aiter_sse():
-                    task_dicts = json.loads(sse.data)
-                    for task_dict in task_dicts:
-                        yield Task.model_validate(task_dict, strict=False)

indexify/executor/task_reporter.py

@@ -1,459 +0,0 @@
-import asyncio
-import hashlib
-import time
-from typing import Any, List, Optional, Tuple
-
-import nanoid
-from httpx import Timeout
-from tensorlake.function_executor.proto.function_executor_pb2 import FunctionOutput
-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,
-    OutputEncoding,
-    ReportTaskOutcomeRequest,
-    TaskOutcome,
-)
-from indexify.proto.executor_api_pb2_grpc import ExecutorAPIStub
-
-from .api_objects import (
-    TASK_OUTCOME_FAILURE,
-    TASK_OUTCOME_SUCCESS,
-    DataPayload,
-    IngestFnOutputsResponse,
-    RouterOutput,
-    TaskResult,
-)
-from .blob_store.blob_store import BLOBStore
-from .function_executor.task_output import TaskOutput
-from .grpc.channel_manager import ChannelManager
-from .metrics.task_reporter import (
-    metric_report_task_outcome_errors,
-    metric_report_task_outcome_latency,
-    metric_report_task_outcome_rpcs,
-    metric_server_ingest_files_errors,
-    metric_server_ingest_files_latency,
-    metric_server_ingest_files_requests,
-    metric_task_output_blob_store_upload_errors,
-    metric_task_output_blob_store_upload_latency,
-    metric_task_output_blob_store_uploads,
-)
-
-
-# https://github.com/psf/requests/issues/1081#issuecomment-428504128
-class ForceMultipartDict(dict):
-    def __bool__(self):
-        return True
-
-
-FORCE_MULTIPART = ForceMultipartDict()
-UTF_8_CONTENT_TYPE = "application/octet-stream"
-
-
-class TaskOutputSummary:
-    def __init__(self):
-        self.output_count: int = 0
-        self.output_total_bytes: int = 0
-        self.router_output_count: int = 0
-        self.stdout_count: int = 0
-        self.stdout_total_bytes: int = 0
-        self.stderr_count: int = 0
-        self.stderr_total_bytes: int = 0
-        self.total_bytes: int = 0
-
-
-class TaskReporter:
-    def __init__(
-        self,
-        base_url: str,
-        executor_id: str,
-        channel_manager: ChannelManager,
-        blob_store: BLOBStore,
-        config_path: Optional[str] = None,
-    ):
-        self._base_url = base_url
-        self._executor_id = executor_id
-        self._is_shutdown = False
-        # Use thread-safe sync client due to issues with async client.
-        # Async client attempts to use connections it already closed.
-        # See e.g. https://github.com/encode/httpx/issues/2337.
-        # Creating a new async client for each request fixes this but it
-        # results in not reusing established TCP connections to server.
-        self._client = get_httpx_client(config_path, make_async=False)
-        self._channel_manager = channel_manager
-        self._blob_store = blob_store
-
-    async def shutdown(self) -> None:
-        """Shuts down the task reporter.
-
-        Task reporter stops reporting all task outcomes to the Server.
-        There are many task failures due to Executor shutdown. We give wrong
-        signals to Server if we report such failures.
-        """
-        self._is_shutdown = True
-
-    async def report(self, output: TaskOutput, logger: Any) -> None:
-        """Reports result of the supplied task."""
-        logger = logger.bind(module=__name__)
-
-        if self._is_shutdown:
-            logger.warning(
-                "task reporter got shutdown, skipping task outcome reporting"
-            )
-            return
-
-        # TODO: If the files are uploaded successfully,
-        # we should record that so that if we fail to report
-        # the task outcome, we don't retry the upload.
-        # This will save us some time and resources.
-        # It's good to do this once we delete all the legacy code paths.
-
-        output_summary: TaskOutputSummary = _task_output_summary(output)
-        logger.info(
-            "reporting task outcome",
-            total_bytes=output_summary.total_bytes,
-            total_files=output_summary.output_count
-            + output_summary.stdout_count
-            + output_summary.stderr_count,
-            output_files=output_summary.output_count,
-            output_bytes=output_summary.total_bytes,
-            router_output_count=output_summary.router_output_count,
-            stdout_bytes=output_summary.stdout_total_bytes,
-            stderr_bytes=output_summary.stderr_total_bytes,
-        )
-
-        if output.output_payload_uri_prefix is None:
-            ingested_files = await self._ingest_files_at_server(output, logger)
-        else:
-            ingested_files = await self._ingest_files_at_blob_store(output, logger)
-
-        fn_outputs = []
-        for data_payload in ingested_files.data_payloads:
-            fn_outputs.append(
-                DataPayloadProto(
-                    path=data_payload.path,  # TODO: stop using this deprecated field once Server side migration is done.
-                    uri=data_payload.path,
-                    size=data_payload.size,
-                    sha256_hash=data_payload.sha256_hash,
-                    encoding=_to_grpc_data_payload_encoding(output),
-                    encoding_version=0,
-                )
-            )
-        stdout, stderr = None, None
-        if ingested_files.stdout is not None:
-            stdout = DataPayloadProto(
-                path=ingested_files.stdout.path,  # TODO: stop using this deprecated field once Server side migration is done.
-                uri=ingested_files.stdout.path,
-                size=ingested_files.stdout.size,
-                sha256_hash=ingested_files.stdout.sha256_hash,
-                encoding=DataPayloadEncoding.DATA_PAYLOAD_ENCODING_UTF8_TEXT,
-                encoding_version=0,
-            )
-        if ingested_files.stderr is not None:
-            stderr = DataPayloadProto(
-                path=ingested_files.stderr.path,  # TODO: stop using this deprecated field once Server side migration is done.
-                uri=ingested_files.stderr.path,
-                size=ingested_files.stderr.size,
-                sha256_hash=ingested_files.stderr.sha256_hash,
-                encoding=DataPayloadEncoding.DATA_PAYLOAD_ENCODING_UTF8_TEXT,
-                encoding_version=0,
-            )
-
-        request = ReportTaskOutcomeRequest(
-            task_id=output.task_id,
-            namespace=output.namespace,
-            graph_name=output.graph_name,
-            function_name=output.function_name,
-            graph_invocation_id=output.graph_invocation_id,
-            outcome=_to_grpc_task_outcome(output),
-            invocation_id=output.graph_invocation_id,
-            executor_id=self._executor_id,
-            reducer=output.reducer,
-            next_functions=(output.router_output.edges if output.router_output else []),
-            fn_outputs=fn_outputs,
-            stdout=stdout,
-            stderr=stderr,
-            output_encoding=_to_grpc_output_encoding(output),
-            output_encoding_version=0,
-        )
-        try:
-            stub = ExecutorAPIStub(await self._channel_manager.get_channel())
-            with (
-                metric_report_task_outcome_latency.time(),
-                metric_report_task_outcome_errors.count_exceptions(),
-            ):
-                metric_report_task_outcome_rpcs.inc()
-                await stub.report_task_outcome(request, timeout=5.0)
-        except Exception as e:
-            logger.error("failed to report task outcome", error=e)
-            raise e
-
-    async def _ingest_files_at_server(
-        self, output: TaskOutput, logger: Any
-    ) -> IngestFnOutputsResponse:
-        logger.warning("uploading task output files to server (deprecated mode)")
-
-        task_result, output_files = self._process_task_output(output)
-        task_result_data = task_result.model_dump_json(exclude_none=True)
-
-        kwargs = {
-            "data": {"task_result": task_result_data},
-            # Use httpx default timeout of 5s for all timeout types.
-            # For read timeouts, use 5 minutes to allow for large file uploads.
-            "timeout": Timeout(
-                5.0,
-                read=5.0 * 60,
-            ),
-            "files": output_files if len(output_files) > 0 else FORCE_MULTIPART,
-        }
-
-        start_time = time.time()
-        with metric_server_ingest_files_latency.time():
-            metric_server_ingest_files_requests.inc()
-            # Run in a separate thread to not block the main event loop.
-            response = await asyncio.to_thread(
-                self._client.post,
-                url=f"{self._base_url}/internal/ingest_fn_outputs",
-                **kwargs,
-            )
-        end_time = time.time()
-        logger.info(
-            "files uploaded to server",
-            response_time=end_time - start_time,
-            response_code=response.status_code,
-        )
-
-        try:
-            response.raise_for_status()
-        except Exception as e:
-            metric_server_ingest_files_errors.inc()
-            # Caller catches and logs the exception.
-            raise Exception(
-                "failed to upload files. "
-                f"Response code: {response.status_code}. "
-                f"Response text: '{response.text}'."
-            ) from e
-
-        ingested_files_response = response.json()
-        return IngestFnOutputsResponse.model_validate(ingested_files_response)
-
-    async def _ingest_files_at_blob_store(
-        self, output: TaskOutput, logger: Any
-    ) -> IngestFnOutputsResponse:
-        start_time = time.time()
-        with (
-            metric_task_output_blob_store_upload_latency.time(),
-            metric_task_output_blob_store_upload_errors.count_exceptions(),
-        ):
-            metric_task_output_blob_store_uploads.inc()
-            response = await self._upload_output_to_blob_store(output, logger)
-
-        logger.info(
-            "files uploaded to blob store",
-            duration=time.time() - start_time,
-        )
-        return response
-
-    async def _upload_output_to_blob_store(
-        self, output: TaskOutput, logger: Any
-    ) -> IngestFnOutputsResponse:
-        data_payloads: List[DataPayload] = []
-        stdout: Optional[DataPayload] = None
-        stderr: Optional[DataPayload] = None
-
-        if output.stdout is not None:
-            stdout_url = f"{output.output_payload_uri_prefix}.{output.task_id}.stdout"
-            stdout_bytes: bytes = output.stdout.encode()
-            await self._blob_store.put(stdout_url, stdout_bytes, logger)
-            stdout = DataPayload(
-                path=stdout_url,
-                size=len(stdout_bytes),
-                sha256_hash=_compute_hash(stdout_bytes),
-            )
-
-        if output.stderr is not None:
-            stderr_url = f"{output.output_payload_uri_prefix}.{output.task_id}.stderr"
-            stderr_bytes: bytes = output.stderr.encode()
-            await self._blob_store.put(stderr_url, stderr_bytes, logger)
-            stderr = DataPayload(
-                path=stderr_url,
-                size=len(stderr_bytes),
-                sha256_hash=_compute_hash(stderr_bytes),
-            )
-
-        if output.function_output is not None:
-            for func_output_item in output.function_output.outputs:
-                node_output_sequence = len(data_payloads)
-                if output.reducer:
-                    # Reducer tasks have to write their results into the same blob.
-                    output_url = (
-                        f"{output.output_payload_uri_prefix}.{node_output_sequence}"
-                    )
-                else:
-                    # Regular tasks write their results into different blobs made unique using task ids.
-                    output_url = f"{output.output_payload_uri_prefix}.{output.task_id}.{node_output_sequence}"
-
-                output_bytes: bytes = (
-                    func_output_item.bytes
-                    if func_output_item.HasField("bytes")
-                    else func_output_item.string.encode()
-                )
-                await self._blob_store.put(output_url, output_bytes, logger)
-                data_payloads.append(
-                    DataPayload(
-                        path=output_url,
-                        size=len(output_bytes),
-                        sha256_hash=_compute_hash(output_bytes),
-                    )
-                )
-
-        return IngestFnOutputsResponse(
-            data_payloads=data_payloads,
-            stdout=stdout,
-            stderr=stderr,
-        )
-
-    def _process_task_output(self, output: TaskOutput) -> Tuple[TaskResult, List[Any]]:
-        task_result = TaskResult(
-            outcome="failure",
-            namespace=output.namespace,
-            compute_graph=output.graph_name,
-            compute_fn=output.function_name,
-            invocation_id=output.graph_invocation_id,
-            executor_id=self._executor_id,
-            task_id=output.task_id,
-            reducer=output.reducer,
-        )
-        output_files: List[Any] = []
-        task_result.outcome = (
-            TASK_OUTCOME_SUCCESS if output.success else TASK_OUTCOME_FAILURE
-        )
-
-        _process_function_output(
-            function_output=output.function_output, output_files=output_files
-        )
-        _process_router_output(
-            router_output=output.router_output, task_result=task_result
-        )
-        _process_stdout(stdout=output.stdout, output_files=output_files)
-        _process_stderr(stderr=output.stderr, output_files=output_files)
-
-        return task_result, output_files
-
-
-def _process_function_output(
-    function_output: Optional[FunctionOutput], output_files: List[Any]
-) -> None:
-    if function_output is None:
-        return
-
-    for output in function_output.outputs or []:
-        payload = output.bytes if output.HasField("bytes") else output.string
-        output_files.append(
-            (
-                "node_outputs",
-                (nanoid.generate(), payload, output.content_type),
-            )
-        )
-
-
-def _process_router_output(
-    router_output: Optional[RouterOutput],
-    task_result: TaskResult,
-) -> None:
-    if router_output is None:
-        return
-
-    task_result.router_output = RouterOutput(edges=router_output.edges)
-
-
-def _process_stdout(stdout: Optional[str], output_files: List[Any]) -> None:
-    if stdout is None:
-        return
-
-    output_files.append(
-        (
-            "stdout",
-            (
-                nanoid.generate(),
-                stdout.encode(),
-                UTF_8_CONTENT_TYPE,
-            ),
-        )
-    )
-
-
-def _process_stderr(stderr: Optional[str], output_files: List[Any]) -> None:
-    if stderr is None:
-        return
-
-    output_files.append(
-        (
-            "stderr",
-            (
-                nanoid.generate(),
-                stderr.encode(),
-                UTF_8_CONTENT_TYPE,
-            ),
-        )
-    )
-
-
-def _task_output_summary(output: TaskOutput) -> TaskOutputSummary:
-    summary: TaskOutputSummary = TaskOutputSummary()
-
-    if output.stdout is not None:
-        summary.stdout_count += 1
-        summary.stdout_total_bytes += len(output.stdout)
-
-    if output.stderr is not None:
-        summary.stderr_count += 1
-        summary.stderr_total_bytes += len(output.stderr)
-
-    if output.function_output is not None:
-        for func_output_item in output.function_output.outputs:
-            output_len: bytes = len(
-                func_output_item.bytes
-                if func_output_item.HasField("bytes")
-                else func_output_item.string
-            )
-            summary.output_count += 1
-            summary.output_total_bytes += output_len
-
-    if output.router_output is not None:
-        summary.router_output_count += 1
-
-    summary.total_bytes = (
-        summary.output_total_bytes
-        + summary.stdout_total_bytes
-        + summary.stderr_total_bytes
-    )
-    return summary
-
-
-def _to_grpc_task_outcome(task_output: TaskOutput) -> TaskOutcome:
-    if task_output.success:
-        return TaskOutcome.TASK_OUTCOME_SUCCESS
-    else:
-        return TaskOutcome.TASK_OUTCOME_FAILURE
-
-
-def _to_grpc_output_encoding(task_output: TaskOutput) -> OutputEncoding:
-    if task_output.output_encoding == "json":
-        return OutputEncoding.OUTPUT_ENCODING_JSON
-    else:
-        return OutputEncoding.OUTPUT_ENCODING_PICKLE
-
-
-def _to_grpc_data_payload_encoding(task_output: TaskOutput) -> DataPayloadEncoding:
-    if task_output.output_encoding == "json":
-        return DataPayloadEncoding.DATA_PAYLOAD_ENCODING_UTF8_JSON
-    else:
-        return DataPayloadEncoding.DATA_PAYLOAD_ENCODING_BINARY_PICKLE
-
-
-def _compute_hash(data: bytes) -> str:
-    hasher = hashlib.sha256(usedforsecurity=False)
-    hasher.update(data)
-    return hasher.hexdigest()