indexify 0.4.9__py3-none-any.whl → 0.4.11__py3-none-any.whl

indexify/executor/executor.py

@@ -52,7 +52,7 @@ class Executor:
         blob_store: BLOBStore,
         host_resources_provider: HostResourcesProvider,
     ):
-        self._logger = structlog.get_logger(module=__name__)
+        self._logger = structlog.get_logger(module=__name__, executor_id=id)
         protocol: str = "http"
         if config_path:
             self._logger.info("running the extractor with TLS enabled")

indexify/executor/function_executor/function_executor.py

@@ -7,6 +7,8 @@ import grpc
 from tensorlake.function_executor.proto.function_executor_pb2 import (
     InfoRequest,
     InfoResponse,
+    InitializationFailureReason,
+    InitializationOutcomeCode,
     InitializeRequest,
     InitializeResponse,
 )
@@ -315,16 +317,42 @@ async def _initialize_server(
                 initialize_request,
                 timeout=customer_code_timeout_sec,
             )
-            # TODO: set real stdout and stderr when their proper capturing on FE initialization is implemented.
-            if initialize_response.success:
-                return FunctionExecutorInitializationResult()
-            elif initialize_response.HasField("customer_error"):
+
+            if (
+                initialize_response.outcome_code
+                == InitializationOutcomeCode.INITIALIZE_OUTCOME_CODE_SUCCESS
+            ):
                 return FunctionExecutorInitializationResult(
-                    error=FunctionExecutorInitializationError.FUNCTION_ERROR,
-                    stderr=initialize_response.customer_error,
+                    stdout=initialize_response.stdout, stderr=initialize_response.stderr
                 )
+            elif (
+                initialize_response.outcome_code
+                == InitializationOutcomeCode.INITIALIZE_OUTCOME_CODE_FAILURE
+            ):
+                if (
+                    initialize_response.failure_reason
+                    == InitializationFailureReason.INITIALIZATION_FAILURE_REASON_FUNCTION_ERROR
+                ):
+                    return FunctionExecutorInitializationResult(
+                        error=FunctionExecutorInitializationError.FUNCTION_ERROR,
+                        stdout=initialize_response.stdout,
+                        stderr=initialize_response.stderr,
+                    )
+                elif (
+                    initialize_response.failure_reason
+                    == InitializationFailureReason.INITIALIZATION_FAILURE_REASON_INTERNAL_ERROR
+                ):
+                    # Don't add stdout/stderr because this is customer data.
+                    raise RuntimeError("initialize RPC failed with internal error")
+                else:
+                    raise ValueError(
+                        f"unexpected failure reason {InitializationFailureReason.Name(initialize_response.failure_reason)} in initialize RPC response"
+                    )
             else:
-                raise Exception("initialize RPC failed at function executor server")
+                raise ValueError(
+                    f"unexpected outcome code {InitializationOutcomeCode.Name(initialize_response.outcome_code)} in initialize RPC response"
+                )
+
         except grpc.aio.AioRpcError as e:
             if e.code() == grpc.StatusCode.DEADLINE_EXCEEDED:
                 return FunctionExecutorInitializationResult(

indexify/executor/function_executor/invocation_state_client.py

@@ -8,6 +8,7 @@ from tensorlake.function_executor.proto.function_executor_pb2 import (
     InvocationStateRequest,
     InvocationStateResponse,
     SerializedObject,
+    SerializedObjectEncoding,
     SetInvocationStateResponse,
 )
 from tensorlake.function_executor.proto.function_executor_pb2_grpc import (
@@ -25,6 +26,10 @@ from .metrics.invocation_state_client import (
     metric_server_set_state_requests,
 )
 
+# We're currently only supporting CloudPickle for invocation state values.
+# FIXME: if Function Executor sends us something else then we fail the calls.
+_VALUE_CONTENT_TYPE = "application/octet-stream"
+
 
 class InvocationStateClient:
     """InvocationStateClient is a client for the invocation state server of a Function Executor.
@@ -196,14 +201,21 @@ class InvocationStateClient:
         url: str = (
             f"{self._base_url}/internal/namespaces/{self._namespace}/compute_graphs/{self._graph}/invocations/{invocation_id}/ctx/{key}"
         )
-        payload = value.bytes if value.HasField("bytes") else value.string
+        if (
+            value.encoding
+            != SerializedObjectEncoding.SERIALIZED_OBJECT_ENCODING_BINARY_PICKLE
+        ):
+            raise ValueError(
+                f"Unsupported value encoding: {SerializedObjectEncoding.Name(value.encoding)}. "
+                "Only binary pickle is supported for invocation state values."
+            )
 
         response = await self._http_client.post(
             url=url,
             files=[
                 (
                     "value",
-                    ("value", payload, value.content_type),
+                    ("value", value.data, _VALUE_CONTENT_TYPE),
                 ),
             ],
         )
@@ -245,7 +257,7 @@ class InvocationStateClient:
             )
             raise
 
-        return serialized_object_from_http_response(response)
+        return _serialized_object_from_http_response(response)
 
     def _validate_request(self, request: InvocationStateRequest) -> None:
         (
@@ -265,17 +277,15 @@ class InvocationStateClient:
             raise ValueError("unknown request type")
 
 
-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_http_response(response: httpx.Response) -> SerializedObject:
+    if response.headers["content-type"] != _VALUE_CONTENT_TYPE:
+        raise ValueError(
+            f"Unexpected content type: {response.headers['content-type']}. "
+            f"Expected: {_VALUE_CONTENT_TYPE}."
         )
+
+    return SerializedObject(
+        data=response.content,
+        encoding=SerializedObjectEncoding.SERIALIZED_OBJECT_ENCODING_BINARY_PICKLE,
+        encoding_version=0,
+    )

indexify/executor/function_executor_controller/completed_task_metrics.py

@@ -66,6 +66,7 @@ def emit_completed_task_metrics(task_info: TaskInfo, logger: Any) -> None:
         elif task_failure_reason in [
             TaskFailureReason.TASK_FAILURE_REASON_FUNCTION_ERROR,
             TaskFailureReason.TASK_FAILURE_REASON_FUNCTION_TIMEOUT,
+            TaskFailureReason.TASK_FAILURE_REASON_INVOCATION_ERROR,
         ]:
             metric_tasks_completed.labels(
                 outcome_code=METRIC_TASKS_COMPLETED_OUTCOME_CODE_FAILURE,

indexify/executor/function_executor_controller/create_function_executor.py

@@ -5,6 +5,7 @@ from typing import Any, Optional, Tuple
 from tensorlake.function_executor.proto.function_executor_pb2 import (
     InitializeRequest,
     SerializedObject,
+    SerializedObjectEncoding,
 )
 
 from indexify.executor.blob_store.blob_store import BLOBStore

indexify/executor/function_executor_controller/destroy_function_executor.py

@@ -1,28 +1,31 @@
+import asyncio
 from typing import Any, Optional
 
 from indexify.executor.function_executor.function_executor import FunctionExecutor
-from indexify.proto.executor_api_pb2 import FunctionExecutorTerminationReason
 
 from .events import FunctionExecutorDestroyed
 
 
 async def destroy_function_executor(
     function_executor: Optional[FunctionExecutor],
-    termination_reason: FunctionExecutorTerminationReason,
+    lock: asyncio.Lock,
     logger: Any,
 ) -> FunctionExecutorDestroyed:
-    """Destroys a function executor if it's not None.
+    """Destroys the function executor if it's not None.
+
+    The supplied lock is used to ensure that if a destroy operation is in progress,
+    then another caller won't return immediately assuming that the destroy is complete
+    due to its idempotency.
 
     Doesn't raise any exceptions.
     """
     logger = logger.bind(module=__name__)
 
     if function_executor is not None:
-        logger.info(
-            "destroying function executor",
-        )
-        await function_executor.destroy()
-
-    return FunctionExecutorDestroyed(
-        is_success=True, termination_reason=termination_reason
-    )
+        async with lock:
+            logger.info(
+                "destroying function executor",
+            )
+            await function_executor.destroy()
+
+    return FunctionExecutorDestroyed(is_success=True)

indexify/executor/function_executor_controller/downloads.py

@@ -4,7 +4,10 @@ from pathlib import Path
 from typing import Any, Optional
 
 import nanoid
-from tensorlake.function_executor.proto.function_executor_pb2 import SerializedObject
+from tensorlake.function_executor.proto.function_executor_pb2 import (
+    SerializedObject,
+    SerializedObjectEncoding,
+)
 
 from indexify.executor.blob_store.blob_store import BLOBStore
 from indexify.proto.executor_api_pb2 import (
@@ -179,20 +182,28 @@ def _serialized_object_from_data_payload_proto(
     """
     if data_payload.encoding == DataPayloadEncoding.DATA_PAYLOAD_ENCODING_BINARY_PICKLE:
         return SerializedObject(
-            bytes=data,
-            content_type="application/octet-stream",
+            data=data,
+            encoding=SerializedObjectEncoding.SERIALIZED_OBJECT_ENCODING_BINARY_PICKLE,
+            encoding_version=data_payload.encoding_version,
         )
     elif data_payload.encoding == DataPayloadEncoding.DATA_PAYLOAD_ENCODING_UTF8_TEXT:
         return SerializedObject(
-            content_type="text/plain",
-            string=data.decode("utf-8"),
+            data=data,
+            encoding=SerializedObjectEncoding.SERIALIZED_OBJECT_ENCODING_UTF8_TEXT,
+            encoding_version=data_payload.encoding_version,
         )
     elif data_payload.encoding == DataPayloadEncoding.DATA_PAYLOAD_ENCODING_UTF8_JSON:
-        result = SerializedObject(
-            content_type="application/json",
-            string=data.decode("utf-8"),
+        return SerializedObject(
+            data=data,
+            encoding=SerializedObjectEncoding.SERIALIZED_OBJECT_ENCODING_UTF8_JSON,
+            encoding_version=data_payload.encoding_version,
+        )
+    elif data_payload.encoding == DataPayloadEncoding.DATA_PAYLOAD_ENCODING_BINARY_ZIP:
+        return SerializedObject(
+            data=data,
+            encoding=SerializedObjectEncoding.SERIALIZED_OBJECT_ENCODING_BINARY_ZIP,
+            encoding_version=data_payload.encoding_version,
         )
-        return result
 
     raise ValueError(
         f"Can't convert data payload {data_payload} into serialized object"

indexify/executor/function_executor_controller/events.py

@@ -55,19 +55,12 @@ class FunctionExecutorDestroyed(BaseEvent):
     Event indicating that Function Executor has been destroyed.
     """
 
-    def __init__(
-        self, is_success: bool, termination_reason: FunctionExecutorTerminationReason
-    ):
+    def __init__(self, is_success: bool):
         super().__init__(EventType.FUNCTION_EXECUTOR_DESTROYED)
         self.is_success: bool = is_success
-        self.termination_reason: FunctionExecutorTerminationReason = termination_reason
 
     def __str__(self) -> str:
-        return (
-            f"Event(type={self.event_type.name}, "
-            f"is_success={self.is_success}, "
-            f"termination_reason={FunctionExecutorTerminationReason.Name(self.termination_reason)})"
-        )
+        return f"Event(type={self.event_type.name}, " f"is_success={self.is_success})"
 
 
 class ShutdownInitiated(BaseEvent):
@@ -75,15 +68,8 @@ class ShutdownInitiated(BaseEvent):
     Event indicating that Function Executor shutdown has been initiated.
     """
 
-    def __init__(self, termination_reason: FunctionExecutorTerminationReason):
+    def __init__(self):
         super().__init__(EventType.SHUTDOWN_INITIATED)
-        self.termination_reason: FunctionExecutorTerminationReason = termination_reason
-
-    def __str__(self) -> str:
-        return (
-            f"Event(type={self.event_type.name}, "
-            f"termination_reason={FunctionExecutorTerminationReason.Name(self.termination_reason)})"
-        )
 
 
 class TaskPreparationFinished(BaseEvent):