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

flyte/_internal/controllers/remote/_controller.py

@@ -0,0 +1,361 @@
+from __future__ import annotations
+
+import asyncio
+from collections import defaultdict
+from pathlib import Path
+from typing import Any, 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._datastructures import ActionID, NativeInterface, SerializationContext
+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
+
+R = TypeVar("R")
+
+
+async def upload_inputs_with_retry(inputs: convert.Inputs, inputs_uri: str) -> None:
+    """
+    Upload inputs to the specified URI with error handling.
+
+    Args:
+        inputs: The inputs to upload
+        inputs_uri: The destination URI
+
+    Raises:
+        RuntimeSystemError: If the upload fails
+    """
+    try:
+        # TODO Add retry decorator to this
+        await io.upload_inputs(inputs, inputs_uri)
+    except Exception as e:
+        logger.exception("Failed to upload inputs", e)
+        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)
+
+
+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 = 100,
+    ):
+        """ """
+        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)
+        )
+
+    async def _submit(self, _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
+
+        inputs = await convert.convert_from_native_to_inputs(_task.native_interface, *args, **kwargs)
+        sub_action_id, sub_action_output_path = convert.generate_sub_action_id_and_output_path(tctx, _task.name, inputs)
+
+        # 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 {sub_action_id.name}")
+            code_bundle = await build_pkl_bundle(
+                _task,
+                upload_to_controlplane=False,
+                upload_from_dataplane_path=io.pkl_path(sub_action_output_path),
+            )
+
+        inputs_uri = io.inputs_path(sub_action_output_path)
+        await upload_inputs_with_retry(inputs, inputs_uri)
+
+        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)
+
+        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,
+        )
+
+        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
+        async with self._parent_action_semaphore[current_action_id.name]:
+            return await self._submit(_task, *args, **kwargs)
+
+    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(action_id.name, None)
+
+    async def get_action_outputs(
+        self, _interface: NativeInterface, _func_name: str, *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_name: 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
+
+        inputs = await convert.convert_from_native_to_inputs(_interface, *args, **kwargs)
+        sub_action_id, sub_action_output_path = convert.generate_sub_action_id_and_output_path(tctx, _func_name, inputs)
+
+        inputs_uri = io.inputs_path(sub_action_output_path)
+        await upload_inputs_with_retry(inputs, inputs_uri)
+
+        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
+
+        native_interface = types.guess_interface(_task.spec.task_template.interface)
+        inputs = await convert.convert_from_native_to_inputs(native_interface, *args, **kwargs)
+        sub_action_id, sub_action_output_path = convert.generate_sub_action_id_and_output_path(tctx, task_name, inputs)
+
+        inputs_uri = io.inputs_path(sub_action_output_path)
+        await upload_inputs_with_retry(inputs, inputs_uri)
+
+        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,
+        )
+
+        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