indexify 0.3.20__tar.gz → 0.3.21__tar.gz

{indexify-0.3.20 → indexify-0.3.21}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.3
 Name: indexify
-Version: 0.3.20
+Version: 0.3.21
 Summary: Open Source Indexify components and helper tools
 Home-page: https://github.com/tensorlakeai/indexify
 License: Apache 2.0

{indexify-0.3.20 → indexify-0.3.21}/pyproject.toml

@@ -1,7 +1,7 @@
 [tool.poetry]
 name = "indexify"
 # Incremented if any of the components provided in this packages are updated.
-version = "0.3.20"
+version = "0.3.21"
 description = "Open Source Indexify components and helper tools"
 authors = ["Tensorlake Inc. <support@tensorlake.ai>"]
 license = "Apache 2.0"

{indexify-0.3.20 → indexify-0.3.21}/src/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-0.3.20 → indexify-0.3.21}/src/indexify/executor/blob_store/local_fs_blob_store.py

@@ -27,7 +27,7 @@ class LocalFSBLOBStore:
 
     def _sync_get(self, path: str) -> bytes:
         if not os.path.isabs(path):
-            raise ValueError(f"Path {path} is not absolute")
+            raise ValueError(f"Path {path} must be absolute")
 
         if os.path.exists(path):
             with open(path, mode="rb") as blob_file:
@@ -37,7 +37,7 @@ class LocalFSBLOBStore:
 
     def _sync_put(self, path: str, value: bytes) -> None:
         if not os.path.isabs(path):
-            raise ValueError(f"Path {path} is not absolute")
+            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:

{indexify-0.3.20 → indexify-0.3.21}/src/indexify/executor/blob_store/s3_blob_store.py

@@ -77,8 +77,11 @@ class S3BLOBStore:
 
 
 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}'")

{indexify-0.3.20 → indexify-0.3.21}/src/indexify/executor/downloader.py

@@ -1,6 +1,6 @@
 import asyncio
 import os
-from typing import Any, Optional
+from typing import Any, Optional, Union
 
 import httpx
 import nanoid
@@ -8,11 +8,12 @@ from tensorlake.function_executor.proto.function_executor_pb2 import SerializedO
 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 (
-    DataPayload,
     DataPayloadEncoding,
 )
 
+from .api_objects import DataPayload
 from .blob_store.blob_store import BLOBStore
 from .metrics.downloader import (
     metric_graph_download_errors,
@@ -49,7 +50,7 @@ class Downloader:
         namespace: str,
         graph_name: str,
         graph_version: str,
-        data_payload: Optional[DataPayload],
+        data_payload: Optional[Union[DataPayload, DataPayloadProto]],
         logger: Any,
     ) -> SerializedObject:
         logger = logger.bind(module=__name__)
@@ -99,7 +100,7 @@ class Downloader:
         function_name: str,
         graph_invocation_id: str,
         reducer_output_key: str,
-        data_payload: Optional[DataPayload],
+        data_payload: Optional[Union[DataPayload, DataPayloadProto]],
         logger: Any,
     ) -> SerializedObject:
         logger = logger.bind(module=__name__)
@@ -124,7 +125,7 @@ class Downloader:
         namespace: str,
         graph_name: str,
         graph_version: str,
-        data_payload: Optional[DataPayload],
+        data_payload: Optional[Union[DataPayload, DataPayloadProto]],
         logger: Any,
     ) -> SerializedObject:
         # Cache graph to reduce load on the server.
@@ -151,7 +152,7 @@ class Downloader:
                 graph_version=graph_version,
                 logger=logger,
             )
-        else:
+        elif isinstance(data_payload, DataPayloadProto):
             (
                 MessageValidator(data_payload)
                 .required_field("uri")
@@ -160,7 +161,15 @@ class Downloader:
             data: bytes = await self._blob_store.get(
                 uri=data_payload.uri, logger=logger
             )
-            return _data_payload_to_serialized_object(
+            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,
             )
@@ -204,7 +213,7 @@ class Downloader:
         graph_name: str,
         graph_invocation_id: str,
         input_key: str,
-        data_payload: Optional[DataPayload],
+        data_payload: Optional[Union[DataPayload, DataPayloadProto]],
         logger: Any,
     ) -> SerializedObject:
         if data_payload is None:
@@ -221,7 +230,7 @@ class Downloader:
                 return await self._fetch_function_input_from_server(
                     input_key=input_key, logger=logger
                 )
-        else:
+        elif isinstance(data_payload, DataPayloadProto):
             (
                 MessageValidator(data_payload)
                 .required_field("uri")
@@ -230,7 +239,15 @@ class Downloader:
             data: bytes = await self._blob_store.get(
                 uri=data_payload.uri, logger=logger
             )
-            return _data_payload_to_serialized_object(
+            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,
             )
@@ -242,7 +259,7 @@ class Downloader:
         function_name: str,
         graph_invocation_id: str,
         reducer_output_key: str,
-        data_payload: Optional[DataPayload],
+        data_payload: Optional[Union[DataPayload, DataPayloadProto]],
         logger: Any,
     ) -> SerializedObject:
         if data_payload is None:
@@ -254,7 +271,7 @@ class Downloader:
                 reducer_output_key=reducer_output_key,
                 logger=logger,
             )
-        else:
+        elif isinstance(data_payload, DataPayloadProto):
             (
                 MessageValidator(data_payload)
                 .required_field("uri")
@@ -263,7 +280,15 @@ class Downloader:
             data: bytes = await self._blob_store.get(
                 uri=data_payload.uri, logger=logger
             )
-            return _data_payload_to_serialized_object(
+            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,
             )
@@ -315,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()
@@ -346,8 +375,23 @@ def serialized_object_from_http_response(response: httpx.Response) -> Serialized
         )
 
 
-def _data_payload_to_serialized_object(
+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.
 

{indexify-0.3.20 → indexify-0.3.21}/src/indexify/executor/executor.py

@@ -128,6 +128,7 @@ class Executor:
             executor_id=id,
             config_path=config_path,
             channel_manager=self._channel_manager,
+            blob_store=blob_store,
         )
 
         # HTTP mode task runner
@@ -261,6 +262,7 @@ class Executor:
                 function_name=task.compute_fn,
                 graph_version=task.graph_version,
                 graph_invocation_id=task.invocation_id,
+                output_payload_uri_prefix=task.output_payload_uri_prefix,
             )
             logger.error("task execution failed", exc_info=e)
 
@@ -304,19 +306,19 @@ class Executor:
             graph_name=task.compute_graph,
             graph_version=task.graph_version,
             logger=logger,
-            data_payload=None,
+            data_payload=task.graph_payload,
         )
         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,
+            data_payload=task.input_payload,
             logger=logger,
         )
         init_value: Optional[SerializedObject] = (
             None
-            if task.reducer_output_id is None
+            if task.reducer_output_id is None and task.reducer_input_payload is None
             else (
                 await self._downloader.download_init_value(
                     namespace=task.namespace,
@@ -324,7 +326,7 @@ class Executor:
                     function_name=task.compute_fn,
                     graph_invocation_id=task.invocation_id,
                     reducer_output_key=task.reducer_output_id,
-                    data_payload=None,
+                    data_payload=task.reducer_input_payload,
                     logger=logger,
                 )
             )

{indexify-0.3.20 → indexify-0.3.21}/src/indexify/executor/function_executor/single_task_runner.py

@@ -96,6 +96,7 @@ class SingleTaskRunner:
                     graph_invocation_id=self._task_input.task.invocation_id,
                     stderr=str(e),
                     success=False,
+                    output_payload_uri_prefix=self._task_input.task.output_payload_uri_prefix,
                 )
 
         try:
@@ -311,6 +312,7 @@ def _task_output(task: Task, response: RunTaskResponse) -> TaskOutput:
         reducer=response.is_reducer,
         success=response.success,
         metrics=metrics,
+        output_payload_uri_prefix=task.output_payload_uri_prefix,
     )
 
     if response.HasField("function_output"):

{indexify-0.3.20 → indexify-0.3.21}/src/indexify/executor/function_executor/task_output.py

@@ -25,6 +25,7 @@ class TaskOutput:
         function_name: str,
         graph_version: str,
         graph_invocation_id: str,
+        output_payload_uri_prefix: Optional[str],
         output_encoding: Optional[str] = None,
         function_output: Optional[FunctionOutput] = None,
         router_output: Optional[RouterOutput] = None,
@@ -50,6 +51,7 @@ class TaskOutput:
         self.is_internal_error = is_internal_error
         self.metrics = metrics
         self.output_encoding = output_encoding
+        self.output_payload_uri_prefix = output_payload_uri_prefix
 
     @classmethod
     def internal_error(
@@ -60,6 +62,7 @@ class TaskOutput:
         function_name: str,
         graph_version: str,
         graph_invocation_id: str,
+        output_payload_uri_prefix: Optional[str],
     ) -> "TaskOutput":
         """Creates a TaskOutput for an internal error."""
         # We are not sharing internal error messages with the customer.
@@ -72,6 +75,7 @@ class TaskOutput:
             graph_invocation_id=graph_invocation_id,
             stderr="Platform failed to execute the function.",
             is_internal_error=True,
+            output_payload_uri_prefix=output_payload_uri_prefix,
         )
 
     @classmethod
@@ -84,6 +88,7 @@ class TaskOutput:
         graph_version: str,
         graph_invocation_id: str,
         timeout_sec: float,
+        output_payload_uri_prefix: Optional[str],
     ) -> "TaskOutput":
         """Creates a TaskOutput for an function timeout error."""
         # Task stdout, stderr is not available.
@@ -96,4 +101,5 @@ class TaskOutput:
             graph_invocation_id=graph_invocation_id,
             stderr=f"Function exceeded its configured timeout of {timeout_sec:.3f} sec.",
             is_internal_error=False,
+            output_payload_uri_prefix=output_payload_uri_prefix,
         )

{indexify-0.3.20 → indexify-0.3.21}/src/indexify/executor/grpc/task_controller.py

@@ -396,6 +396,11 @@ class TaskController:
             function_name=self._task.function_name,
             graph_version=self._task.graph_version,
             graph_invocation_id=self._task.graph_invocation_id,
+            output_payload_uri_prefix=(
+                self._task.output_payload_uri_prefix
+                if self._task.HasField("output_payload_uri_prefix")
+                else None
+            ),
         )
 
     def _function_timeout_output(self, timeout_sec: float) -> TaskOutput:
@@ -407,6 +412,11 @@ class TaskController:
             graph_version=self._task.graph_version,
             graph_invocation_id=self._task.graph_invocation_id,
             timeout_sec=timeout_sec,
+            output_payload_uri_prefix=(
+                self._task.output_payload_uri_prefix
+                if self._task.HasField("output_payload_uri_prefix")
+                else None
+            ),
         )
 
 
@@ -437,6 +447,11 @@ def _task_output_from_function_executor_response(
         reducer=response.is_reducer,
         success=response.success,
         metrics=metrics,
+        output_payload_uri_prefix=(
+            task.output_payload_uri_prefix
+            if task.HasField("output_payload_uri_prefix")
+            else None
+        ),
     )
 
     if response.HasField("function_output"):

{indexify-0.3.20 → indexify-0.3.21}/src/indexify/executor/metrics/task_reporter.py

@@ -21,6 +21,23 @@ metric_server_ingest_files_latency: prometheus_client.Histogram = (
     )
 )
 
+metric_task_output_blob_store_uploads: prometheus_client.Counter = (
+    prometheus_client.Counter(
+        "task_output_blob_store_uploads", "Number of task output uploads to blob store"
+    )
+)
+metric_task_output_blob_store_upload_errors: prometheus_client.Counter = (
+    prometheus_client.Counter(
+        "task_output_blob_store_upload_errors",
+        "Number of failed task output uploads to blob store",
+    )
+)
+metric_task_output_blob_store_upload_latency: prometheus_client.Histogram = (
+    latency_metric_for_fast_operation(
+        "task_output_blob_store_upload", "Upload task output to blob store"
+    )
+)
+
 metric_report_task_outcome_rpcs = prometheus_client.Counter(
     "report_task_outcome_rpcs",
     "Number of report task outcome RPCs to Server",

{indexify-0.3.20 → indexify-0.3.21}/src/indexify/executor/task_reporter.py

@@ -1,4 +1,5 @@
 import asyncio
+import hashlib
 import time
 from typing import Any, List, Optional, Tuple
 
@@ -7,8 +8,8 @@ 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 (
-    DataPayload,
     DataPayloadEncoding,
     OutputEncoding,
     ReportTaskOutcomeRequest,
@@ -19,10 +20,12 @@ 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 (
@@ -32,6 +35,9 @@ from .metrics.task_reporter import (
     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,
 )
 
 
@@ -63,6 +69,7 @@ class TaskReporter:
         base_url: str,
         executor_id: str,
         channel_manager: ChannelManager,
+        blob_store: BLOBStore,
         config_path: Optional[str] = None,
     ):
         self._base_url = base_url
@@ -75,6 +82,7 @@ class TaskReporter:
         # 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.
@@ -95,9 +103,13 @@ class TaskReporter:
             )
             return
 
-        task_result, output_files, output_summary = self._process_task_output(output)
-        task_result_data = task_result.model_dump_json(exclude_none=True)
+        # 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,
@@ -111,56 +123,15 @@ class TaskReporter:
             stderr_bytes=output_summary.stderr_total_bytes,
         )
 
-        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,
-        }
-
-        # TODO: Instead of uploading the files to server, upload them to S3.
-        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",
-            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
-
-        # 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.
+        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)
 
-        ingested_files_response = response.json()
-        ingested_files = IngestFnOutputsResponse.model_validate(ingested_files_response)
         fn_outputs = []
         for data_payload in ingested_files.data_payloads:
             fn_outputs.append(
-                DataPayload(
+                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,
@@ -170,8 +141,8 @@ class TaskReporter:
                 )
             )
         stdout, stderr = None, None
-        if ingested_files.stdout:
-            stdout = DataPayload(
+        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,
@@ -179,8 +150,8 @@ class TaskReporter:
                 encoding=DataPayloadEncoding.DATA_PAYLOAD_ENCODING_UTF8_TEXT,
                 encoding_version=0,
             )
-        if ingested_files.stderr:
-            stderr = DataPayload(
+        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,
@@ -218,9 +189,132 @@ class TaskReporter:
             logger.error("failed to report task outcome", error=e)
             raise e
 
-    def _process_task_output(
-        self, output: TaskOutput
-    ) -> Tuple[TaskResult, List[Any], TaskOutputSummary]:
+    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,
@@ -231,9 +325,8 @@ class TaskReporter:
             task_id=output.task_id,
         )
         output_files: List[Any] = []
-        summary: TaskOutputSummary = TaskOutputSummary()
         if output is None:
-            return task_result, output_files, summary
+            return task_result, output_files
 
         task_result.outcome = (
             TASK_OUTCOME_SUCCESS if output.success else TASK_OUTCOME_FAILURE
@@ -241,33 +334,19 @@ class TaskReporter:
         task_result.reducer = output.reducer
 
         _process_function_output(
-            function_output=output.function_output,
-            output_files=output_files,
-            summary=summary,
+            function_output=output.function_output, output_files=output_files
         )
         _process_router_output(
-            router_output=output.router_output, task_result=task_result, summary=summary
-        )
-        _process_stdout(
-            stdout=output.stdout, output_files=output_files, summary=summary
-        )
-        _process_stderr(
-            stderr=output.stderr, output_files=output_files, summary=summary
+            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)
 
-        summary.total_bytes = (
-            summary.output_total_bytes
-            + summary.stdout_total_bytes
-            + summary.stderr_total_bytes
-        )
-
-        return task_result, output_files, summary
+        return task_result, output_files
 
 
 def _process_function_output(
-    function_output: Optional[FunctionOutput],
-    output_files: List[Any],
-    summary: TaskOutputSummary,
+    function_output: Optional[FunctionOutput], output_files: List[Any]
 ) -> None:
     if function_output is None:
         return
@@ -280,25 +359,19 @@ def _process_function_output(
                 (nanoid.generate(), payload, output.content_type),
             )
         )
-        summary.output_count += 1
-        summary.output_total_bytes += len(payload)
 
 
 def _process_router_output(
     router_output: Optional[RouterOutput],
     task_result: TaskResult,
-    summary: TaskOutputSummary,
 ) -> None:
     if router_output is None:
         return
 
     task_result.router_output = RouterOutput(edges=router_output.edges)
-    summary.router_output_count += 1
 
 
-def _process_stdout(
-    stdout: Optional[str], output_files: List[Any], summary: TaskOutputSummary
-) -> None:
+def _process_stdout(stdout: Optional[str], output_files: List[Any]) -> None:
     if stdout is None:
         return
 
@@ -312,13 +385,9 @@ def _process_stdout(
             ),
         )
     )
-    summary.stdout_count += 1
-    summary.stdout_total_bytes += len(stdout)
 
 
-def _process_stderr(
-    stderr: Optional[str], output_files: List[Any], summary: TaskOutputSummary
-) -> None:
+def _process_stderr(stderr: Optional[str], output_files: List[Any]) -> None:
     if stderr is None:
         return
 
@@ -332,8 +401,38 @@ def _process_stderr(
             ),
         )
     )
-    summary.stderr_count += 1
-    summary.stderr_total_bytes += len(stderr)
+
+
+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:
@@ -355,3 +454,9 @@ def _to_grpc_data_payload_encoding(task_output: TaskOutput) -> DataPayloadEncodi
         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()

{indexify-0.3.20 → indexify-0.3.21}/src/indexify/executor/task_runner.py

@@ -85,6 +85,7 @@ class TaskRunner:
                 function_name=task_input.task.compute_fn,
                 graph_version=task_input.task.graph_version,
                 graph_invocation_id=task_input.task.invocation_id,
+                output_payload_uri_prefix=task_input.task.output_payload_uri_prefix,
             )
         finally:
             if state is not None: