flyte 2.0.0b32__py3-none-any.whl

flyte/_interface.py

@@ -0,0 +1,119 @@
+from __future__ import annotations
+
+import inspect
+import typing
+from enum import Enum
+from typing import Dict, Generator, Literal, Tuple, Type, TypeVar, Union, cast, get_args, get_origin, get_type_hints
+
+from flyte._logging import logger
+
+LITERAL_ENUM = "LiteralEnum"
+
+
+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.
+    """
+    if isinstance(return_annotation, str):
+        raise TypeError("String return annotations are not supported.")
+
+    # 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"):
+            # Task returns named tuple
+            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
+        # Task returns unnamed typing.Tuple
+        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)
+        annotations = {}
+        for i, r in enumerate(ra):
+            if r is Ellipsis:
+                raise TypeError("Variable length tuples are not supported as return types.")
+            if get_origin(r) is Literal:
+                annotations[default_output_name(i)] = literal_to_enum(cast(Type, r))
+            else:
+                annotations[default_output_name(i)] = r
+        return annotations
+
+    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
+        # Task returns unnamed native tuple
+        if get_origin(return_annotation) is Literal:
+            return {default_output_name(): literal_to_enum(cast(Type, return_annotation))}
+        return {default_output_name(): cast(Type, return_annotation)}
+
+
+def literal_to_enum(literal_type: Type) -> Type[Enum | typing.Any]:
+    """Convert a Literal[...] into Union[str, Enum]."""
+
+    if get_origin(literal_type) is not Literal:
+        raise TypeError(f"{literal_type} is not a Literal")
+
+    values = get_args(literal_type)
+    if not all(isinstance(v, str) for v in values):
+        logger.warning(f"Literal type {literal_type} contains non-string values, using Any instead of Enum")
+        return typing.Any
+    # Deduplicate & keep order
+    enum_dict = {str(v).upper(): v for v in values}
+
+    # Dynamically create an Enum
+    literal_enum = Enum(LITERAL_ENUM, enum_dict)  # type: ignore
+
+    return literal_enum  # type: ignore

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,129 @@
+import concurrent.futures
+import threading
+from typing import TYPE_CHECKING, Any, Callable, Literal, Optional, Protocol, Tuple, TypeVar
+
+from flyte._task import TaskTemplate
+from flyte.models import ActionID, NativeInterface
+
+if TYPE_CHECKING:
+    from flyte.remote._task import TaskDetails
+
+from ._trace import TraceInfo
+
+__all__ = ["Controller", "ControllerType", "TraceInfo", "create_controller", "get_controller"]
+
+if TYPE_CHECKING:
+    import concurrent.futures
+
+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.
+        """
+        ...
+
+    def submit_sync(self, _task: TaskTemplate, *args, **kwargs) -> concurrent.futures.Future:
+        """
+        This should call the async submit method above, but return a concurrent Future object that can be
+        used in a blocking wait or wrapped in an async future. This is called when
+          a) a synchronous task is kicked off locally,
+          b) a running task (of either kind) kicks off a downstream synchronous task.
+        """
+        ...
+
+    async def submit_task_ref(self, _task: "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: Callable, *args, **kwargs
+    ) -> Tuple[TraceInfo, bool]:
+        """
+        This method returns the outputs of the action, if it is available.
+        :param _interface: NativeInterface
+        :param _func: 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()
+
+
+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 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,239 @@
+import asyncio
+import atexit
+import concurrent.futures
+import os
+import pathlib
+import threading
+from typing import Any, Callable, Tuple, TypeVar
+
+import flyte.errors
+from flyte._cache.cache import VersionParameters, cache_from_request
+from flyte._cache.local_cache import LocalTaskCache
+from flyte._context import internal_ctx
+from flyte._internal.controllers import TraceInfo
+from flyte._internal.runtime import convert
+from flyte._internal.runtime.entrypoints import direct_dispatch
+from flyte._internal.runtime.types_serde import transform_native_to_typed_interface
+from flyte._logging import log, logger
+from flyte._task import AsyncFunctionTaskTemplate, TaskTemplate
+from flyte._utils.helpers import _selector_policy
+from flyte.models import ActionID, NativeInterface
+from flyte.remote._task import TaskDetails
+
+R = TypeVar("R")
+
+
+class _TaskRunner:
+    """A task runner that runs an asyncio event loop on a background thread."""
+
+    def __init__(self) -> None:
+        self.__loop: asyncio.AbstractEventLoop | None = None
+        self.__runner_thread: threading.Thread | None = None
+        self.__lock = threading.Lock()
+        atexit.register(self._close)
+
+    def _close(self) -> None:
+        if self.__loop:
+            self.__loop.stop()
+
+    def _execute(self) -> None:
+        loop = self.__loop
+        assert loop is not None
+        try:
+            loop.run_forever()
+        finally:
+            loop.close()
+
+    def get_exc_handler(self):
+        def exc_handler(loop, context):
+            logger.error(
+                f"Taskrunner for {self.__runner_thread.name if self.__runner_thread else 'no thread'} caught"
+                f" exception in {loop}: {context}"
+            )
+
+        return exc_handler
+
+    def get_run_future(self, coro: Any) -> concurrent.futures.Future:
+        """Synchronously run a coroutine on a background thread."""
+        name = f"{threading.current_thread().name} : loop-runner"
+        with self.__lock:
+            if self.__loop is None:
+                with _selector_policy():
+                    self.__loop = asyncio.new_event_loop()
+
+                exc_handler = self.get_exc_handler()
+                self.__loop.set_exception_handler(exc_handler)
+                self.__runner_thread = threading.Thread(target=self._execute, daemon=True, name=name)
+                self.__runner_thread.start()
+        fut = asyncio.run_coroutine_threadsafe(coro, self.__loop)
+        return fut
+
+
+class LocalController:
+    def __init__(self):
+        logger.debug("LocalController init")
+        self._runner_map: dict[str, _TaskRunner] = {}
+
+    @log
+    async def submit(self, _task: TaskTemplate, *args, **kwargs) -> Any:
+        """
+        Main entrypoint for submitting a task to the local 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)
+        inputs_hash = convert.generate_inputs_hash_from_proto(inputs.proto_inputs)
+        task_interface = transform_native_to_typed_interface(_task.interface)
+
+        sub_action_id, sub_action_output_path = convert.generate_sub_action_id_and_output_path(
+            tctx, _task.name, inputs_hash, 0
+        )
+        sub_action_raw_data_path = tctx.raw_data_path
+        # Make sure the output path exists
+        pathlib.Path(sub_action_output_path).mkdir(parents=True, exist_ok=True)
+        pathlib.Path(sub_action_raw_data_path.path).mkdir(parents=True, exist_ok=True)
+
+        task_cache = cache_from_request(_task.cache)
+        cache_enabled = task_cache.is_enabled()
+        if isinstance(_task, AsyncFunctionTaskTemplate):
+            version_parameters = VersionParameters(func=_task.func, image=_task.image)
+        else:
+            version_parameters = VersionParameters(func=None, image=_task.image)
+        cache_version = task_cache.get_version(version_parameters)
+        cache_key = convert.generate_cache_key_hash(
+            _task.name,
+            inputs_hash,
+            task_interface,
+            cache_version,
+            list(task_cache.get_ignored_inputs()),
+            inputs.proto_inputs,
+        )
+
+        out = None
+        # We only get output from cache if the cache behavior is set to auto
+        if task_cache.behavior == "auto":
+            out = await LocalTaskCache.get(cache_key)
+            if out is not None:
+                logger.info(
+                    f"Cache hit for task '{_task.name}' (version: {cache_version}), getting result from cache..."
+                )
+
+        if out is None:
+            out, err = await direct_dispatch(
+                _task,
+                controller=self,
+                action=sub_action_id,
+                raw_data_path=sub_action_raw_data_path,
+                inputs=inputs,
+                version=cache_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")
+
+            # store into cache
+            if cache_enabled and out is not None:
+                await LocalTaskCache.set(cache_key, out)
+
+        if _task.native_interface.outputs:
+            if out is None:
+                raise flyte.errors.RuntimeSystemError("BadOutput", "Task output not captured.")
+            result = await convert.convert_outputs_to_native(_task.native_interface, out)
+            return result
+        return None
+
+    def submit_sync(self, _task: TaskTemplate, *args, **kwargs) -> concurrent.futures.Future:
+        name = threading.current_thread().name + f"PID:{os.getpid()}"
+        coro = self.submit(_task, *args, **kwargs)
+        if name not in self._runner_map:
+            if len(self._runner_map) > 100:
+                logger.warning(
+                    "More than 100 event loop runners created!!! This could be a case of runaway recursion..."
+                )
+            self._runner_map[name] = _TaskRunner()
+
+        return self._runner_map[name].get_run_future(coro)
+
+    async def finalize_parent_action(self, action: ActionID):
+        pass
+
+    async def stop(self):
+        await LocalTaskCache.close()
+
+    async def watch_for_errors(self):
+        pass
+
+    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  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
+
+        inputs_hash = convert.generate_inputs_hash_from_proto(converted_inputs.proto_inputs)
+        action_id, action_output_path = convert.generate_sub_action_id_and_output_path(
+            tctx,
+            _func.__name__,
+            inputs_hash,
+            0,
+        )
+        assert action_output_path
+        return (
+            TraceInfo(
+                name=_func.__name__,
+                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, info.name)
+            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.start_time
+        assert info.end_time
+
+    async def submit_task_ref(self, _task: TaskDetails, max_inline_io_bytes: int, *args, **kwargs) -> Any:
+        raise flyte.errors.ReferenceTaskError(
+            f"Reference tasks cannot be executed locally, only remotely. Found remote task {_task.name}"
+        )

flyte/_internal/controllers/_trace.py

@@ -0,0 +1,48 @@
+from dataclasses import dataclass, field
+from typing import Any, Optional
+
+from flyteidl2.core import interface_pb2
+
+from flyte.models 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.
+    """
+
+    name: str
+    action: ActionID
+    interface: NativeInterface
+    inputs_path: str
+    start_time: float = field(init=False, default=0.0)
+    end_time: float = field(init=False, default=0.0)
+    output: Optional[Any] = None
+    error: Optional[Exception] = None
+    typed_interface: Optional[interface_pb2.TypedInterface] = None
+
+    def add_outputs(self, output: Any, start_time: float, end_time: float):
+        """
+        Add outputs to the trace information.
+        :param output: Output of the action
+        :param start_time: Start time of the action
+        :param end_time: End time of the action
+        :return:
+        """
+        self.output = output
+        self.start_time = start_time
+        self.end_time = end_time
+
+    def add_error(self, error: Exception, start_time: float, end_time: float):
+        """
+        Add error to the trace information.
+        :param error: Error of the action
+        :param start_time: Start time of the action
+        :param end_time: End time of the action
+        :return:
+        """
+        self.error = error
+        self.start_time = start_time
+        self.end_time = end_time

flyte/_internal/controllers/remote/__init__.py

@@ -0,0 +1,58 @@
+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,
+    endpoint: str | None = None,
+    insecure: bool = False,
+    insecure_skip_verify: bool = False,
+    ca_cert_file_path: str | None = None,
+    client_config: ClientConfig | None = None,
+    auth_type: AuthType = "Pkce",
+    headless: bool = False,
+    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.
+    """
+    assert endpoint or api_key, "Either endpoint or api_key must be provided when initializing remote controller"
+    from ._client import ControllerClient
+    from ._controller import RemoteController
+
+    if endpoint:
+        client_coro = ControllerClient.for_endpoint(
+            endpoint,
+            insecure=insecure,
+            insecure_skip_verify=insecure_skip_verify,
+            ca_cert_file_path=ca_cert_file_path,
+            client_id=client_id,
+            client_credentials_secret=client_credentials_secret,
+            auth_type=auth_type,
+        )
+    elif api_key:
+        client_coro = ControllerClient.for_api_key(
+            api_key,
+            insecure=insecure,
+            insecure_skip_verify=insecure_skip_verify,
+            ca_cert_file_path=ca_cert_file_path,
+            client_id=client_id,
+            client_credentials_secret=client_credentials_secret,
+            auth_type=auth_type,
+        )
+
+    controller = RemoteController(
+        client_coro=client_coro,
+    )
+    return controller