flyte 0.1.0__py3-none-any.whl → 0.2.0a0__py3-none-any.whl

flyte/_internal/controllers/remote/_controller.py

@@ -0,0 +1,494 @@
+from __future__ import annotations
+
+import asyncio
+import concurrent.futures
+import os
+import threading
+from collections import defaultdict
+from collections.abc import Callable
+from pathlib import Path
+from typing import Any, AsyncIterable, Awaitable, DefaultDict, Tuple, TypeVar
+
+import flyte
+import flyte.errors
+import flyte.storage as storage
+import flyte.types as types
+from flyte._code_bundle import build_pkl_bundle
+from flyte._context import internal_ctx
+from flyte._internal.controllers import TraceInfo
+from flyte._internal.controllers.remote._action import Action
+from flyte._internal.controllers.remote._core import Controller
+from flyte._internal.controllers.remote._service_protocol import ClientSet
+from flyte._internal.runtime import convert, io
+from flyte._internal.runtime.task_serde import translate_task_to_wire
+from flyte._logging import logger
+from flyte._protos.workflow import run_definition_pb2, task_definition_pb2
+from flyte._task import TaskTemplate
+from flyte._utils.helpers import _selector_policy
+from flyte.models import ActionID, NativeInterface, SerializationContext
+
+R = TypeVar("R")
+
+
+async def upload_inputs_with_retry(serialized_inputs: AsyncIterable[bytes] | bytes, inputs_uri: str) -> None:
+    """
+    Upload inputs to the specified URI with error handling.
+
+    Args:
+        serialized_inputs: The serialized inputs to upload
+        inputs_uri: The destination URI
+
+    Raises:
+        RuntimeSystemError: If the upload fails
+    """
+    try:
+        # TODO Add retry decorator to this
+        await storage.put_stream(serialized_inputs, to_path=inputs_uri)
+    except Exception as e:
+        logger.exception("Failed to upload inputs")
+        raise flyte.errors.RuntimeSystemError(type(e).__name__, str(e)) from e
+
+
+async def handle_action_failure(action: Action, task_name: str) -> Exception:
+    """
+    Handle action failure by loading error details or raising a RuntimeSystemError.
+
+    Args:
+        action: The updated action
+        task_name: The name of the task
+
+    Raises:
+        Exception: The converted native exception or RuntimeSystemError
+    """
+    err = action.err or action.client_err
+    if not err and action.phase == run_definition_pb2.PHASE_FAILED:
+        logger.error(f"Server reported failure for action {action.name}, checking error file.")
+        try:
+            error_path = io.error_path(f"{action.run_output_base}/{action.action_id.name}/1")
+            err = await io.load_error(error_path)
+        except Exception as e:
+            logger.exception("Failed to load error file", e)
+            err = flyte.errors.RuntimeSystemError(type(e).__name__, f"Failed to load error file: {e}")
+    else:
+        logger.error(f"Server reported failure for action {action.action_id.name}, error: {err}")
+
+    exc = convert.convert_error_to_native(err)
+    if not exc:
+        return flyte.errors.RuntimeSystemError("UnableToConvertError", f"Error in task {task_name}: {err}")
+    return exc
+
+
+async def load_and_convert_outputs(iface: NativeInterface, realized_outputs_uri: str) -> Any:
+    """
+    Load outputs from the given URI and convert them to native format.
+
+    Args:
+        iface: The Native interface
+        realized_outputs_uri: The URI where outputs are stored
+
+    Returns:
+        The converted native outputs
+    """
+    outputs_file_path = io.outputs_path(realized_outputs_uri)
+    outputs = await io.load_outputs(outputs_file_path)
+    return await convert.convert_outputs_to_native(iface, outputs)
+
+
+def unique_action_name(action_id: ActionID) -> str:
+    return f"{action_id.name}_{action_id.run_name}"
+
+
+class RemoteController(Controller):
+    """
+    This a specialized controller that wraps the core controller and performs IO, serialization and deserialization
+    """
+
+    def __init__(
+        self,
+        client_coro: Awaitable[ClientSet],
+        workers: int,
+        max_system_retries: int,
+        default_parent_concurrency: int = 1000,
+    ):
+        """ """
+        super().__init__(
+            client_coro=client_coro,
+            workers=workers,
+            max_system_retries=max_system_retries,
+        )
+        self._default_parent_concurrency = default_parent_concurrency
+        self._parent_action_semaphore: DefaultDict[str, asyncio.Semaphore] = defaultdict(
+            lambda: asyncio.Semaphore(default_parent_concurrency)
+        )
+        self._parent_action_task_call_sequence: DefaultDict[str, DefaultDict[int, int]] = defaultdict(
+            lambda: defaultdict(int)
+        )
+        self._submit_loop: asyncio.AbstractEventLoop | None = None
+        self._submit_thread: threading.Thread | None = None
+
+    def generate_task_call_sequence(self, task_obj: object, action_id: ActionID) -> int:
+        """
+        Generate a task call sequence for the given task object and action ID.
+        This is used to track the number of times a task is called within an action.
+        """
+        current_action_sequencer = self._parent_action_task_call_sequence[unique_action_name(action_id)]
+        current_task_id = id(task_obj)
+        v = current_action_sequencer[current_task_id]
+        new_seq = v + 1
+        current_action_sequencer[current_task_id] = new_seq
+        return new_seq
+
+    async def _submit(self, _task_call_seq: int, _task: TaskTemplate, *args, **kwargs) -> Any:
+        ctx = internal_ctx()
+        tctx = ctx.data.task_context
+        if tctx is None:
+            raise flyte.errors.RuntimeSystemError("BadContext", "Task context not initialized")
+        current_action_id = tctx.action
+
+        # In the case of a regular code bundle, we will just pass it down as it is to the downstream tasks
+        # It is not allowed to change the code bundle (for regular code bundles) in the middle of a run.
+        code_bundle = tctx.code_bundle
+
+        if code_bundle and code_bundle.pkl:
+            logger.debug(f"Building new pkl bundle for task {_task.name}")
+            code_bundle = await build_pkl_bundle(
+                _task,
+                upload_to_controlplane=False,
+                upload_from_dataplane_base_path=tctx.run_base_dir,
+            )
+
+        inputs = await convert.convert_from_native_to_inputs(_task.native_interface, *args, **kwargs)
+
+        root_dir = Path(code_bundle.destination).absolute() if code_bundle else Path.cwd()
+        # Don't set output path in sec context because node executor will set it
+        new_serialization_context = SerializationContext(
+            project=current_action_id.project,
+            domain=current_action_id.domain,
+            org=current_action_id.org,
+            code_bundle=code_bundle,
+            version=tctx.version,
+            # supplied version.
+            # input_path=inputs_uri,
+            image_cache=tctx.compiled_image_cache,
+            root_dir=root_dir,
+        )
+
+        task_spec = translate_task_to_wire(_task, new_serialization_context)
+        serialized_inputs = inputs.proto_inputs.SerializeToString(deterministic=True)
+
+        inputs_hash = convert.generate_inputs_hash(serialized_inputs)
+        sub_action_id, sub_action_output_path = convert.generate_sub_action_id_and_output_path(
+            tctx, task_spec, inputs_hash, _task_call_seq
+        )
+
+        inputs_uri = io.inputs_path(sub_action_output_path)
+        await upload_inputs_with_retry(serialized_inputs, inputs_uri)
+
+        md = task_spec.task_template.metadata
+        ignored_input_vars = []
+        if len(md.cache_ignore_input_vars) > 0:
+            ignored_input_vars = list(md.cache_ignore_input_vars)
+        cache_key = None
+        if task_spec.task_template.metadata and task_spec.task_template.metadata.discoverable:
+            discovery_version = task_spec.task_template.metadata.discovery_version
+            cache_key = convert.generate_cache_key_hash(
+                _task.name,
+                inputs_hash,
+                task_spec.task_template.interface,
+                discovery_version,
+                ignored_input_vars,
+                inputs.proto_inputs,
+            )
+
+        # Clear to free memory
+        serialized_inputs = None  # type: ignore
+        inputs_hash = None  # type: ignore
+
+        action = Action.from_task(
+            sub_action_id=run_definition_pb2.ActionIdentifier(
+                name=sub_action_id.name,
+                run=run_definition_pb2.RunIdentifier(
+                    name=current_action_id.run_name,
+                    project=current_action_id.project,
+                    domain=current_action_id.domain,
+                    org=current_action_id.org,
+                ),
+            ),
+            parent_action_name=current_action_id.name,
+            group_data=tctx.group_data,
+            task_spec=task_spec,
+            inputs_uri=inputs_uri,
+            run_output_base=tctx.run_base_dir,
+            cache_key=cache_key,
+        )
+
+        try:
+            logger.info(
+                f"Submitting action Run:[{action.run_name}, Parent:[{action.parent_action_name}], "
+                f"task:[{_task.name}], action:[{action.name}]"
+            )
+            n = await self.submit_action(action)
+            logger.info(f"Action for task [{_task.name}] action id: {action.name}, completed!")
+        except asyncio.CancelledError:
+            # If the action is cancelled, we need to cancel the action on the server as well
+            logger.info(f"Action {action.action_id.name} cancelled, cancelling on server")
+            await self.cancel_action(action)
+            raise
+
+        if n.has_error() or n.phase == run_definition_pb2.PHASE_FAILED:
+            exc = await handle_action_failure(action, _task.name)
+            raise exc
+
+        if _task.native_interface.outputs:
+            if not n.realized_outputs_uri:
+                raise flyte.errors.RuntimeSystemError(
+                    "RuntimeError",
+                    f"Task {n.action_id.name} did not return an output path, but the task has outputs defined.",
+                )
+            return await load_and_convert_outputs(_task.native_interface, n.realized_outputs_uri)
+        return None
+
+    async def submit(self, _task: TaskTemplate, *args, **kwargs) -> Any:
+        """
+        Submit a task to the remote controller.This creates a new action on the queue service.
+        """
+        ctx = internal_ctx()
+        tctx = ctx.data.task_context
+        if tctx is None:
+            raise flyte.errors.RuntimeSystemError("BadContext", "Task context not initialized")
+        current_action_id = tctx.action
+        task_call_seq = self.generate_task_call_sequence(_task, current_action_id)
+        async with self._parent_action_semaphore[unique_action_name(current_action_id)]:
+            return await self._submit(task_call_seq, _task, *args, **kwargs)
+
+    def _sync_thread_loop_runner(self) -> None:
+        """This method runs the event loop and should be invoked in a separate thread."""
+
+        loop = self._submit_loop
+        assert loop is not None
+        try:
+            loop.run_forever()
+        finally:
+            loop.close()
+
+    def submit_sync(self, _task: TaskTemplate, *args, **kwargs) -> concurrent.futures.Future:
+        """
+        This function creates a cached thread and loop for the purpose of calling the submit method synchronously,
+        returning a concurrent Future that can be awaited. There's no need for a lock because this function itself is
+        single threaded and non-async. This pattern here is basically the trivial/degenerate case of the thread pool
+        in the LocalController.
+        Please see additional comments in protocol.
+
+        :param _task:
+        :param args:
+        :param kwargs:
+        :return:
+        """
+        if self._submit_thread is None:
+            # Please see LocalController for the general implementation of this pattern.
+            def exc_handler(loop, context):
+                logger.error(f"Remote controller submit sync loop caught exception in {loop}: {context}")
+
+            with _selector_policy():
+                self._submit_loop = asyncio.new_event_loop()
+                self._submit_loop.set_exception_handler(exc_handler)
+
+            self._submit_thread = threading.Thread(
+                name=f"remote-controller-{os.getpid()}-submitter", daemon=True, target=self._sync_thread_loop_runner
+            )
+            self._submit_thread.start()
+
+        coro = self.submit(_task, *args, **kwargs)
+        assert self._submit_loop is not None, "Submit loop should always have been initialized by now"
+        fut = asyncio.run_coroutine_threadsafe(coro, self._submit_loop)
+        return fut
+
+    async def finalize_parent_action(self, action_id: ActionID):
+        """
+        This method is invoked when the parent action is finished. It will finalize the run and upload the outputs
+        to the control plane.
+        """
+        run_id = run_definition_pb2.RunIdentifier(
+            name=action_id.run_name,
+            project=action_id.project,
+            domain=action_id.domain,
+            org=action_id.org,
+        )
+        await super()._finalize_parent_action(run_id=run_id, parent_action_name=action_id.name)
+        self._parent_action_semaphore.pop(unique_action_name(action_id), None)
+        self._parent_action_task_call_sequence.pop(unique_action_name(action_id), None)
+
+    async def get_action_outputs(
+        self, _interface: NativeInterface, _func: Callable, *args, **kwargs
+    ) -> Tuple[TraceInfo, bool]:
+        """
+        This method returns the outputs of the action, if it is available.
+        If not available it raises a NotFoundError.
+        :param _interface: NativeInterface
+        :param _func: Function name
+        :param args: Arguments
+        :param kwargs: Keyword arguments
+        :return:
+        """
+        ctx = internal_ctx()
+        tctx = ctx.data.task_context
+        if tctx is None:
+            raise flyte.errors.RuntimeSystemError("BadContext", "Task context not initialized")
+        current_action_id = tctx.action
+
+        func_name = _func.__name__
+        invoke_seq_num = self.generate_task_call_sequence(_func, current_action_id)
+        inputs = await convert.convert_from_native_to_inputs(_interface, *args, **kwargs)
+        serialized_inputs = inputs.proto_inputs.SerializeToString(deterministic=True)
+
+        sub_action_id, sub_action_output_path = convert.generate_sub_action_id_and_output_path(
+            tctx, func_name, serialized_inputs, invoke_seq_num
+        )
+
+        inputs_uri = io.inputs_path(sub_action_output_path)
+        await upload_inputs_with_retry(serialized_inputs, inputs_uri)
+        # Clear to free memory
+        serialized_inputs = None  # type: ignore
+
+        prev_action = await self.get_action(
+            run_definition_pb2.ActionIdentifier(
+                name=sub_action_id.name,
+                run=run_definition_pb2.RunIdentifier(
+                    name=current_action_id.run_name,
+                    project=current_action_id.project,
+                    domain=current_action_id.domain,
+                    org=current_action_id.org,
+                ),
+            ),
+            current_action_id.name,
+        )
+
+        if prev_action is None:
+            return TraceInfo(sub_action_id, _interface, inputs_uri), False
+
+        if prev_action.phase == run_definition_pb2.PHASE_FAILED:
+            if prev_action.has_error():
+                exc = convert.convert_error_to_native(prev_action.err)
+                return TraceInfo(sub_action_id, _interface, inputs_uri, error=exc), True
+            else:
+                logger.warning(f"Action {prev_action.action_id.name} failed, but no error was found, re-running trace!")
+        elif prev_action.realized_outputs_uri is not None:
+            outputs_file_path = io.outputs_path(prev_action.realized_outputs_uri)
+            o = await io.load_outputs(outputs_file_path)
+            outputs = await convert.convert_outputs_to_native(_interface, o)
+            return TraceInfo(sub_action_id, _interface, inputs_uri, output=outputs), True
+
+        return TraceInfo(sub_action_id, _interface, inputs_uri), False
+
+    async def record_trace(self, info: TraceInfo):
+        """
+        Record a trace action. This is used to record the trace of the action and should be called when the action
+        :param info:
+        :return:
+        """
+        ctx = internal_ctx()
+        tctx = ctx.data.task_context
+        if tctx is None:
+            raise flyte.errors.RuntimeSystemError("BadContext", "Task context not initialized")
+
+        current_output_path = tctx.output_path
+        sub_run_output_path = storage.join(current_output_path, info.action.name)
+
+        if info.interface.has_outputs():
+            if info.output:
+                outputs = await convert.convert_from_native_to_outputs(info.output, info.interface)
+                outputs_file_path = io.outputs_path(sub_run_output_path)
+                await io.upload_outputs(outputs, outputs_file_path)
+            elif info.error:
+                err = convert.convert_from_native_to_error(info.error)
+                error_path = io.error_path(sub_run_output_path)
+                await io.upload_error(err.err, error_path)
+            else:
+                raise flyte.errors.RuntimeSystemError("BadTraceInfo", "Trace info does not have output or error")
+
+    async def submit_task_ref(self, _task: task_definition_pb2.TaskDetails, *args, **kwargs) -> Any:
+        ctx = internal_ctx()
+        tctx = ctx.data.task_context
+        if tctx is None:
+            raise flyte.errors.RuntimeSystemError("BadContext", "Task context not initialized")
+        current_action_id = tctx.action
+        task_name = _task.spec.task_template.id.name
+
+        invoke_seq_num = self.generate_task_call_sequence(_task, current_action_id)
+
+        native_interface = types.guess_interface(
+            _task.spec.task_template.interface, default_inputs=_task.spec.default_inputs
+        )
+        inputs = await convert.convert_from_native_to_inputs(native_interface, *args, **kwargs)
+        serialized_inputs = inputs.proto_inputs.SerializeToString(deterministic=True)
+        inputs_hash = convert.generate_inputs_hash(serialized_inputs)
+        sub_action_id, sub_action_output_path = convert.generate_sub_action_id_and_output_path(
+            tctx, task_name, inputs_hash, invoke_seq_num
+        )
+
+        inputs_uri = io.inputs_path(sub_action_output_path)
+        await upload_inputs_with_retry(serialized_inputs, inputs_uri)
+        # cache key - task name, task signature, inputs, cache version
+        cache_key = None
+        md = _task.spec.task_template.metadata
+        ignored_input_vars = []
+        if len(md.cache_ignore_input_vars) > 0:
+            ignored_input_vars = list(md.cache_ignore_input_vars)
+        if _task.spec.task_template.metadata and _task.spec.task_template.metadata.discoverable:
+            discovery_version = _task.spec.task_template.metadata.discovery_version
+            cache_key = convert.generate_cache_key_hash(
+                task_name,
+                inputs_hash,
+                _task.spec.task_template.interface,
+                discovery_version,
+                ignored_input_vars,
+                inputs.proto_inputs,
+            )
+
+        # Clear to free memory
+        serialized_inputs = None  # type: ignore
+        inputs_hash = None  # type: ignore
+
+        action = Action.from_task(
+            sub_action_id=run_definition_pb2.ActionIdentifier(
+                name=sub_action_id.name,
+                run=run_definition_pb2.RunIdentifier(
+                    name=current_action_id.run_name,
+                    project=current_action_id.project,
+                    domain=current_action_id.domain,
+                    org=current_action_id.org,
+                ),
+            ),
+            parent_action_name=current_action_id.name,
+            group_data=tctx.group_data,
+            task_spec=_task.spec,
+            inputs_uri=inputs_uri,
+            run_output_base=tctx.run_base_dir,
+            cache_key=cache_key,
+        )
+
+        try:
+            logger.info(
+                f"Submitting action Run:[{action.run_name}, Parent:[{action.parent_action_name}], "
+                f"task:[{task_name}], action:[{action.name}]"
+            )
+            n = await self.submit_action(action)
+            logger.info(f"Action for task [{task_name}] action id: {action.name}, completed!")
+        except asyncio.CancelledError:
+            # If the action is cancelled, we need to cancel the action on the server as well
+            logger.info(f"Action {action.action_id.name} cancelled, cancelling on server")
+            await self.cancel_action(action)
+            raise
+
+        if n.has_error() or n.phase == run_definition_pb2.PHASE_FAILED:
+            exc = await handle_action_failure(action, task_name)
+            raise exc
+
+        if native_interface.outputs:
+            if not n.realized_outputs_uri:
+                raise flyte.errors.RuntimeSystemError(
+                    "RuntimeError",
+                    f"Task {n.action_id.name} did not return an output path, but the task has outputs defined.",
+                )
+            return await load_and_convert_outputs(native_interface, n.realized_outputs_uri)
+        return None