flyte 0.0.1b0__py3-none-any.whl

flyte/_interface.py

@@ -0,0 +1,84 @@
+from __future__ import annotations
+
+import inspect
+from typing import Dict, Generator, Tuple, Type, TypeVar, Union, cast, get_args, get_type_hints
+
+from flyte._logging import logger
+
+
+def default_output_name(index: int = 0) -> str:
+    return f"o{index}"
+
+
+def output_name_generator(length: int) -> Generator[str, None, None]:
+    for x in range(length):
+        yield default_output_name(x)
+
+
+def extract_return_annotation(return_annotation: Union[Type, Tuple, None]) -> Dict[str, Type]:
+    """
+    The input to this function should be sig.return_annotation where sig = inspect.signature(some_func)
+    The purpose of this function is to sort out whether a function is returning one thing, or multiple things, and to
+    name the outputs accordingly, either by using our default name function, or from a typing.NamedTuple.
+
+        # Option 1
+        nt1 = typing.NamedTuple("NT1", x_str=str, y_int=int)
+        def t(a: int, b: str) -> nt1: ...
+
+        # Option 2
+        def t(a: int, b: str) -> typing.NamedTuple("NT1", x_str=str, y_int=int): ...
+
+        # Option 3
+        def t(a: int, b: str) -> typing.Tuple[int, str]: ...
+
+        # Option 4
+        def t(a: int, b: str) -> (int, str): ...
+
+        # Option 5
+        def t(a: int, b: str) -> str: ...
+
+        # Option 6
+        def t(a: int, b: str) -> None: ...
+
+        # Options 7/8
+        def t(a: int, b: str) -> List[int]: ...
+        def t(a: int, b: str) -> Dict[str, int]: ...
+
+    Note that Options 1 and 2 are identical, just syntactic sugar. In the NamedTuple case, we'll use the names in the
+    definition. In all other cases, we'll automatically generate output names, indexed starting at 0.
+    """
+
+    # Handle Option 6
+    # We can think about whether we should add a default output name with type None in the future.
+    if return_annotation in (None, type(None), inspect.Signature.empty):
+        return {}
+
+    # This statement results in true for typing.Namedtuple, single and void return types, so this
+    # handles Options 1, 2. Even though NamedTuple for us is multi-valued, it's a single value for Python
+    if hasattr(return_annotation, "__bases__") and (
+        isinstance(return_annotation, type) or isinstance(return_annotation, TypeVar)
+    ):
+        # isinstance / issubclass does not work for Namedtuple.
+        # Options 1 and 2
+        bases = return_annotation.__bases__  # type: ignore
+        if len(bases) == 1 and bases[0] is tuple and hasattr(return_annotation, "_fields"):
+            logger.debug(f"Task returns named tuple {return_annotation}")
+            return dict(get_type_hints(cast(Type, return_annotation), include_extras=True))
+
+    if hasattr(return_annotation, "__origin__") and return_annotation.__origin__ is tuple:  # type: ignore
+        # Handle option 3
+        logger.debug(f"Task returns unnamed typing.Tuple {return_annotation}")
+        if len(return_annotation.__args__) == 1:  # type: ignore
+            raise TypeError("Tuples should be used to indicate multiple return values, found only one return variable.")
+        ra = get_args(return_annotation)
+        return dict(zip(list(output_name_generator(len(ra))), ra))
+
+    elif isinstance(return_annotation, tuple):
+        if len(return_annotation) == 1:
+            raise TypeError("Please don't use a tuple if you're just returning one thing.")
+        return dict(zip(list(output_name_generator(len(return_annotation))), return_annotation))
+
+    else:
+        # Handle all other single return types
+        logger.debug(f"Task returns unnamed native tuple {return_annotation}")
+        return {default_output_name(): cast(Type, return_annotation)}

flyte/_internal/__init__.py

@@ -0,0 +1,3 @@
+from .controllers import Controller, ControllerType, create_controller
+
+__all__ = ["Controller", "ControllerType", "create_controller"]

flyte/_internal/controllers/__init__.py

@@ -0,0 +1,115 @@
+import threading
+from typing import Any, Literal, Optional, Protocol, Tuple, TypeVar
+
+from flyte._datastructures import ActionID, NativeInterface
+from flyte._task import TaskTemplate
+
+from ._trace import TraceInfo
+
+__all__ = ["Controller", "ControllerType", "TraceInfo", "create_controller", "get_controller"]
+
+from ..._protos.workflow import task_definition_pb2
+
+ControllerType = Literal["local", "remote"]
+
+R = TypeVar("R")
+
+
+class Controller(Protocol):
+    """
+    Controller interface, that is used to execute tasks. The implementation of this interface,
+    can execute tasks in different ways, such as locally, remotely etc.
+    """
+
+    async def submit(self, _task: TaskTemplate, *args, **kwargs) -> Any:
+        """
+        Submit a node to the controller asynchronously and wait for the result. This is async and will block
+        the current coroutine until the result is available.
+        """
+        ...
+
+    async def submit_task_ref(self, _task: task_definition_pb2.TaskDetails, *args, **kwargs) -> Any:
+        """
+        Submit a task reference to the controller asynchronously and wait for the result. This is async and will block
+        the current coroutine until the result is available.
+        """
+        ...
+
+    async def finalize_parent_action(self, action: ActionID):
+        """
+        Finalize the parent action. This can be called to cleanup the action and should be called after the parent
+        task completes
+        :param action: Action ID
+        :return:
+        """
+        ...
+
+    async def watch_for_errors(self): ...
+
+    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.
+        :param _interface: NativeInterface
+        :param _func_name: Function name
+        :param args: Arguments
+        :param kwargs: Keyword arguments
+        :return: TraceInfo object and a boolean indicating if the action was found.
+        if boolean is False, it means the action is not found and the TraceInfo object will have only min info
+        """
+
+    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
+        is completed.
+        :param info: Trace information
+        :return:
+        """
+        ...
+
+    async def stop(self):
+        """
+        Stops the engine and should be called when the engine is no longer needed.
+        """
+        ...
+
+
+# Internal state holder
+class _ControllerState:
+    controller: Optional[Controller] = None
+    lock = threading.Lock()
+
+
+async def get_controller() -> Controller:
+    """
+    Get the controller instance. Raise an error if it has not been created.
+    """
+    if _ControllerState.controller is not None:
+        return _ControllerState.controller
+    raise RuntimeError("Controller is not initialized. Please call get_or_create_controller() first.")
+
+
+def create_controller(
+    ct: ControllerType,
+    **kwargs,
+) -> Controller:
+    """
+    Create a new instance of the controller, based on the kind and the given configuration.
+    """
+    controller: Controller
+    match ct:
+        case "local":
+            from ._local_controller import LocalController
+
+            controller = LocalController()
+        case "remote" | "hybrid":
+            from flyte._internal.controllers.remote import create_remote_controller
+
+            controller = create_remote_controller(**kwargs)
+        case _:
+            raise ValueError(f"{ct} is not a valid controller type.")
+
+    with _ControllerState.lock:
+        _ControllerState.controller = controller
+        return controller

flyte/_internal/controllers/_local_controller.py

@@ -0,0 +1,118 @@
+from typing import Any, Tuple, TypeVar
+
+import flyte.errors
+from flyte._context import internal_ctx
+from flyte._datastructures import ActionID, NativeInterface, RawDataPath
+from flyte._internal.controllers import TraceInfo
+from flyte._internal.runtime import convert
+from flyte._internal.runtime.entrypoints import direct_dispatch
+from flyte._logging import log, logger
+from flyte._protos.workflow import task_definition_pb2
+from flyte._task import TaskTemplate
+
+R = TypeVar("R")
+
+
+class LocalController:
+    def __init__(self):
+        logger.debug("LocalController init")
+
+    @log
+    async def submit(self, _task: TaskTemplate, *args, **kwargs) -> Any:
+        """
+        Submit a node to the controller
+        """
+        ctx = internal_ctx()
+        tctx = ctx.data.task_context
+        if not tctx:
+            raise flyte.errors.RuntimeSystemError("BadContext", "Task context not initialized")
+
+        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)
+        sub_action_raw_data_path = RawDataPath(path=sub_action_output_path)
+
+        out, err = await direct_dispatch(
+            _task,
+            controller=self,
+            action=sub_action_id,
+            raw_data_path=sub_action_raw_data_path,
+            inputs=inputs,
+            version=tctx.version,
+            checkpoints=tctx.checkpoints,
+            code_bundle=tctx.code_bundle,
+            output_path=sub_action_output_path,
+            run_base_dir=tctx.run_base_dir,
+        )
+        if err:
+            exc = convert.convert_error_to_native(err)
+            if exc:
+                raise exc
+            else:
+                raise flyte.errors.RuntimeSystemError("BadError", "Unknown error")
+        if _task.native_interface.outputs and out is not None:
+            result = await convert.convert_outputs_to_native(_task.native_interface, out)
+            return result
+        return out
+
+    async def finalize_parent_action(self, action: ActionID):
+        pass
+
+    async def stop(self):
+        pass
+
+    async def watch_for_errors(self):
+        pass
+
+    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  flyte.errors.ActionNotFoundError.
+        :return:
+        """
+        ctx = internal_ctx()
+        tctx = ctx.data.task_context
+        if not tctx:
+            raise flyte.errors.NotInTaskContextError("BadContext", "Task context not initialized")
+        converted_inputs = convert.Inputs.empty()
+        if _interface.inputs:
+            converted_inputs = await convert.convert_from_native_to_inputs(_interface, *args, **kwargs)
+            assert converted_inputs
+        action_id, action_output_path = convert.generate_sub_action_id_and_output_path(
+            tctx, _func_name, converted_inputs
+        )
+        assert action_output_path
+        return (
+            TraceInfo(
+                action=action_id,
+                interface=_interface,
+                inputs_path=action_output_path,
+            ),
+            True,
+        )
+
+    async def record_trace(self, info: TraceInfo):
+        """
+        This method records the trace of the action.
+        :param info: Trace information
+        :return:
+        """
+        ctx = internal_ctx()
+        tctx = ctx.data.task_context
+        if not tctx:
+            raise flyte.errors.NotInTaskContextError("BadContext", "Task context not initialized")
+
+        if info.interface.outputs and info.output:
+            # If the result is not an AsyncGenerator, convert it directly
+            converted_outputs = await convert.convert_from_native_to_outputs(info.output, info.interface)
+            assert converted_outputs
+        elif info.error:
+            # If there is an error, convert it to a native error
+            converted_error = convert.convert_from_native_to_error(info.error)
+            assert converted_error
+        assert info.action
+        assert info.duration
+
+    async def submit_task_ref(self, _task: task_definition_pb2.TaskDetails, *args, **kwargs) -> Any:
+        raise flyte.errors.ReferenceTaskError("Reference tasks cannot be executed locally, only remotely.")

flyte/_internal/controllers/_trace.py

@@ -0,0 +1,40 @@
+from dataclasses import dataclass
+from datetime import timedelta
+from typing import Any, Optional
+
+from flyte._datastructures import ActionID, NativeInterface
+
+
+@dataclass
+class TraceInfo:
+    """
+    Trace information for the action. This is used to record the trace of the action and should be called when
+     the action is completed.
+    """
+
+    action: ActionID
+    interface: NativeInterface
+    inputs_path: str
+    duration: Optional[timedelta] = None
+    output: Optional[Any] = None
+    error: Optional[Exception] = None
+
+    def add_outputs(self, output: Any, duration: timedelta):
+        """
+        Add outputs to the trace information.
+        :param output: Output of the action
+        :param duration: Duration of the action
+        :return:
+        """
+        self.output = output
+        self.duration = duration
+
+    def add_error(self, error: Exception, duration: timedelta):
+        """
+        Add error to the trace information.
+        :param error: Error of the action
+        :param duration: Duration of the action
+        :return:
+        """
+        self.error = error
+        self.duration = duration

flyte/_internal/controllers/pbhash.py

@@ -0,0 +1,39 @@
+# This is a module that provides hashing utilities for Protobuf objects.
+import base64
+import hashlib
+import json
+
+from google.protobuf import json_format
+from google.protobuf.message import Message
+
+
+def compute_hash(pb: Message) -> bytes:
+    """
+    Computes a deterministic hash in bytes for the Protobuf object.
+    """
+    try:
+        pb_dict = json_format.MessageToDict(pb)
+        # json.dumps with sorted keys to ensure stability
+        stable_json_str = json.dumps(
+            pb_dict, sort_keys=True, separators=(",", ":")
+        )  # separators to ensure no extra spaces
+    except Exception as e:
+        raise ValueError(f"Failed to marshal Protobuf object {pb} to JSON with error: {e}")
+
+    try:
+        # Deterministically hash the JSON object to a byte array. Using SHA-256 for hashing here,
+        # assuming it provides a consistent hash output.
+        hash_obj = hashlib.sha256(stable_json_str.encode("utf-8"))
+    except Exception as e:
+        raise ValueError(f"Failed to hash JSON for Protobuf object {pb} with error: {e}")
+
+    # The digest is guaranteed to be 32 bytes long
+    return hash_obj.digest()
+
+
+def compute_hash_string(pb: Message) -> str:
+    """
+    Computes a deterministic hash in base64 encoded string for the Protobuf object
+    """
+    hash_bytes = compute_hash(pb)
+    return base64.b64encode(hash_bytes).decode("utf-8")

flyte/_internal/controllers/remote/__init__.py

@@ -0,0 +1,40 @@
+from typing import List
+
+from flyte.remote._client.auth import AuthType, ClientConfig
+
+from ._controller import RemoteController
+
+__all__ = ["RemoteController", "create_remote_controller"]
+
+
+def create_remote_controller(
+    *,
+    api_key: str | None = None,
+    auth_type: AuthType = "Pkce",
+    endpoint: str,
+    client_config: ClientConfig | None = None,
+    headless: bool = False,
+    insecure: bool = False,
+    insecure_skip_verify: bool = False,
+    ca_cert_file_path: str | None = None,
+    command: List[str] | None = None,
+    proxy_command: List[str] | None = None,
+    client_id: str | None = None,
+    client_credentials_secret: str | None = None,
+    rpc_retries: int = 3,
+    http_proxy_url: str | None = None,
+) -> RemoteController:
+    """
+    Create a new instance of the remote controller.
+    """
+    from ._client import ControllerClient
+    from ._controller import RemoteController
+
+    controller = RemoteController(
+        client_coro=ControllerClient.for_endpoint(
+            endpoint=endpoint, insecure=insecure, insecure_skip_verify=insecure_skip_verify
+        ),
+        workers=10,
+        max_system_retries=5,
+    )
+    return controller

flyte/_internal/controllers/remote/_action.py

@@ -0,0 +1,141 @@
+from __future__ import annotations
+
+from dataclasses import dataclass
+
+from flyteidl.core import execution_pb2
+
+from flyte._datastructures import GroupData
+from flyte._protos.workflow import run_definition_pb2, state_service_pb2, task_definition_pb2
+
+
+@dataclass
+class Action:
+    """
+    Coroutine safe, as we never do await operations in any method.
+    Holds the inmemory state of a task. It is combined representation of local and remote states.
+    """
+
+    action_id: run_definition_pb2.ActionIdentifier
+    parent_action_name: str
+    friendly_name: str | None = None
+    group: GroupData | None = None
+    task: task_definition_pb2.TaskSpec | None = None
+    inputs_uri: str | None = None
+    run_output_base: str | None = None
+    realized_outputs_uri: str | None = None
+    err: execution_pb2.ExecutionError | None = None
+    phase: run_definition_pb2.Phase | None = None
+    started: bool = False
+    retries: int = 0
+    client_err: Exception | None = None  # This error is set when something goes wrong in the controller.
+
+    @property
+    def name(self) -> str:
+        return self.action_id.name
+
+    @property
+    def run_name(self) -> str:
+        return self.action_id.run.name
+
+    def is_terminal(self) -> bool:
+        """Check if resource has reached terminal state"""
+        if self.phase is None:
+            return False
+        return self.phase in [
+            run_definition_pb2.Phase.PHASE_FAILED,
+            run_definition_pb2.Phase.PHASE_SUCCEEDED,
+            run_definition_pb2.Phase.PHASE_ABORTED,
+            run_definition_pb2.Phase.PHASE_TIMED_OUT,
+        ]
+
+    def increment_retries(self):
+        self.retries += 1
+
+    def is_started(self) -> bool:
+        """Check if resource has been started."""
+        return self.started
+
+    def mark_started(self):
+        self.started = True
+        self.task = None
+
+    def mark_cancelled(self):
+        self.mark_started()
+        self.phase = run_definition_pb2.Phase.PHASE_ABORTED
+
+    def merge_state(self, obj: state_service_pb2.ActionUpdate):
+        """
+        This method is invoked when the watch API sends an update about the state of the action. We need to merge
+        the state of the action with the current state of the action. It is possible that we have no phase information
+        prior to this.
+        :param obj:
+        :return:
+        """
+        if self.phase != obj.phase:
+            self.phase = obj.phase
+            self.err = obj.error if obj.HasField("error") else None
+        self.realized_outputs_uri = obj.output_uri
+        self.started = True
+
+    def merge_in_action_from_submit(self, action: Action):
+        """
+        This method is invoked when parent_action submits an action that was observed previously observed from the
+         watch. We need to merge in the contents of the action, while preserving the observed phase.
+
+        :param action: The submitted action
+        """
+        self.run_output_base = action.run_output_base
+        self.inputs_uri = action.inputs_uri
+        self.group = action.group
+        self.friendly_name = action.friendly_name
+        if not self.started:
+            self.task = action.task
+
+    def set_client_error(self, exc: Exception):
+        self.client_err = exc
+
+    def has_error(self) -> bool:
+        return self.client_err is not None or self.err is not None
+
+    @classmethod
+    def from_task(
+        cls,
+        parent_action_name: str,
+        sub_action_id: run_definition_pb2.ActionIdentifier,
+        group_data: GroupData | None,
+        task_spec: task_definition_pb2.TaskSpec,
+        inputs_uri: str,
+        run_output_base: str,
+    ) -> Action:
+        return cls(
+            action_id=sub_action_id,
+            parent_action_name=parent_action_name,
+            friendly_name=task_spec.task_template.id.name,
+            group=group_data,
+            task=task_spec,
+            inputs_uri=inputs_uri,
+            run_output_base=run_output_base,
+        )
+
+    @classmethod
+    def from_state(cls, parent_action_name: str, obj: state_service_pb2.ActionUpdate) -> Action:
+        """
+        This creates a new action, from the watch api. This is possible in the case of a recovery, where the
+        state service knows about future actions and sends this information to the informer. We may not have
+        encountered the "task" itself yet, but we know about the action id and the state of the action.
+
+        :param parent_action_name:
+        :param obj:
+        :return:
+        """
+        from flyte._logging import logger
+
+        logger.info(f"In Action from_state {obj.action_id} {obj.phase} {obj.output_uri}")
+        return cls(
+            action_id=obj.action_id,
+            parent_action_name=parent_action_name,
+            phase=obj.phase,
+            started=True,
+            err=obj.error if obj.HasField("error") else None,
+            realized_outputs_uri=obj.output_uri,
+        )

flyte/_internal/controllers/remote/_client.py

@@ -0,0 +1,43 @@
+from __future__ import annotations
+
+import grpc.aio
+
+from flyte._protos.workflow import queue_service_pb2_grpc, state_service_pb2_grpc
+from flyte.remote import create_channel
+
+from ._service_protocol import QueueService, StateService
+
+
+class ControllerClient:
+    """
+    A client for the Controller API.
+    """
+
+    def __init__(self, channel: grpc.aio.Channel):
+        self._channel = channel
+        self._state_service = state_service_pb2_grpc.StateServiceStub(channel=channel)
+        self._queue_service = queue_service_pb2_grpc.QueueServiceStub(channel=channel)
+
+    @classmethod
+    async def for_endpoint(cls, endpoint: str, insecure: bool = False, **kwargs) -> ControllerClient:
+        return cls(await create_channel(endpoint, insecure=insecure, **kwargs))
+
+    @property
+    def state_service(self) -> StateService:
+        """
+        The state service.
+        """
+        return self._state_service
+
+    @property
+    def queue_service(self) -> QueueService:
+        """
+        The queue service.
+        """
+        return self._queue_service
+
+    def close(self, grace: float | None = None):
+        """
+        Close the channel.
+        """
+        return self._channel.close(grace=grace)