flyte 0.0.1b0__py3-none-any.whl → 2.0.0b46__py3-none-any.whl

flyte/_interface.py

@@ -1,10 +1,14 @@
 from __future__ import annotations
 
 import inspect
-from typing import Dict, Generator, Tuple, Type, TypeVar, Union, cast, get_args, get_type_hints
+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}"
@@ -47,6 +51,8 @@ def extract_return_annotation(return_annotation: Union[Type, Tuple, None]) -> Di
     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.
@@ -62,16 +68,24 @@ def extract_return_annotation(return_annotation: Union[Type, Tuple, None]) -> Di
         # 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}")
+            # 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
-        logger.debug(f"Task returns unnamed typing.Tuple {return_annotation}")
+        # 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)
-        return dict(zip(list(output_name_generator(len(ra))), ra))
+        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:
@@ -80,5 +94,26 @@ def extract_return_annotation(return_annotation: Union[Type, Tuple, None]) -> Di
 
     else:
         # Handle all other single return types
-        logger.debug(f"Task returns unnamed native tuple {return_annotation}")
+        # 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/controllers/__init__.py

@@ -1,14 +1,19 @@
+import concurrent.futures
 import threading
-from typing import Any, Literal, Optional, Protocol, Tuple, TypeVar
+from typing import TYPE_CHECKING, Any, Callable, Literal, Optional, Protocol, Tuple, TypeVar
 
-from flyte._datastructures import ActionID, NativeInterface
 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"]
 
-from ..._protos.workflow import task_definition_pb2
+if TYPE_CHECKING:
+    import concurrent.futures
 
 ControllerType = Literal["local", "remote"]
 
@@ -28,7 +33,16 @@ class Controller(Protocol):
         """
         ...
 
-    async def submit_task_ref(self, _task: task_definition_pb2.TaskDetails, *args, **kwargs) -> Any:
+    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.
@@ -47,12 +61,12 @@ class Controller(Protocol):
     async def watch_for_errors(self): ...
 
     async def get_action_outputs(
-        self, _interface: NativeInterface, _func_name: str, *args, **kwargs
+        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_name: Function name
+        :param _func: Function name
         :param args: Arguments
         :param kwargs: Keyword arguments
         :return: TraceInfo object and a boolean indicating if the action was found.
@@ -81,13 +95,13 @@ class _ControllerState:
     lock = threading.Lock()
 
 
-async def get_controller() -> Controller:
+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.")
+    raise RuntimeError("Controller is not initialized. Please call create_controller() first.")
 
 
 def create_controller(

flyte/_internal/controllers/_local_controller.py

@@ -1,26 +1,83 @@
-from typing import Any, Tuple, TypeVar
+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._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._internal.runtime.types_serde import transform_native_to_typed_interface
 from flyte._logging import log, logger
-from flyte._protos.workflow import task_definition_pb2
-from flyte._task import TaskTemplate
+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:
         """
-        Submit a node to the controller
+        Main entrypoint for submitting a task to the local controller.
         """
         ctx = internal_ctx()
         tctx = ctx.data.task_context
@@ -28,43 +85,100 @@ class LocalController:
             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,
+        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,
         )
-        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:
+
+        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 out
+        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):
-        pass
+        await LocalTaskCache.close()
 
     async def watch_for_errors(self):
-        pass
+        try:
+            await asyncio.Event().wait()  # Wait indefinitely until cancelled
+        except asyncio.CancelledError:
+            return  # Return with no errors when cancelled
 
     async def get_action_outputs(
-        self, _interface: NativeInterface, _func_name: str, *args, **kwargs
+        self, _interface: NativeInterface, _func: Callable, *args, **kwargs
     ) -> Tuple[TraceInfo, bool]:
         """
         This method returns the outputs of the action, if it is available.
@@ -75,16 +189,23 @@ class LocalController:
         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, converted_inputs
+            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,
@@ -105,14 +226,17 @@ class LocalController:
 
         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)
+            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.duration
+        assert info.start_time
+        assert info.end_time
 
-    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.")
+    async def submit_task_ref(self, _task: TaskDetails, max_inline_io_bytes: int, *args, **kwargs) -> Any:
+        raise flyte.errors.RemoteTaskError(
+            f"Remote tasks cannot be executed locally, only remotely. Found remote task {_task.name}"
+        )

flyte/_internal/controllers/_trace.py

@@ -1,8 +1,9 @@
-from dataclasses import dataclass
-from datetime import timedelta
+from dataclasses import dataclass, field
 from typing import Any, Optional
 
-from flyte._datastructures import ActionID, NativeInterface
+from flyteidl2.core import interface_pb2
+
+from flyte.models import ActionID, NativeInterface
 
 
 @dataclass
@@ -12,29 +13,36 @@ class TraceInfo:
      the action is completed.
     """
 
+    name: str
     action: ActionID
     interface: NativeInterface
     inputs_path: str
-    duration: Optional[timedelta] = None
+    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, duration: timedelta):
+    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 duration: Duration of the action
+        :param start_time: Start time of the action
+        :param end_time: End time of the action
         :return:
         """
         self.output = output
-        self.duration = duration
+        self.start_time = start_time
+        self.end_time = end_time
 
-    def add_error(self, error: Exception, duration: timedelta):
+    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 duration: Duration of the action
+        :param start_time: Start time of the action
+        :param end_time: End time of the action
         :return:
         """
         self.error = error
-        self.duration = duration
+        self.start_time = start_time
+        self.end_time = end_time

flyte/_internal/controllers/remote/__init__.py

@@ -10,13 +10,13 @@ __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,
+    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,
@@ -27,14 +27,43 @@ def create_remote_controller(
     """
     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
 
+    # https://grpc.io/docs/guides/keepalive/#keepalive-configuration-specification
+    channel_options = [
+        ("grpc.keepalive_permit_without_calls", 1),
+        ("grpc.keepalive_time_ms", 30000),  # Send keepalive ping every 30 seconds
+        ("grpc.keepalive_timeout_ms", 10000),  # Wait 10 seconds for keepalive response
+        ("grpc.http2.max_pings_without_data", 0),  # Allow unlimited pings without data
+        ("grpc.http2.min_ping_interval_without_data_ms", 30000),  # Min 30s between pings
+    ]
+
+    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,
+            grpc_options=channel_options,
+        )
+    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,
+            grpc_options=channel_options,
+        )
+
     controller = RemoteController(
-        client_coro=ControllerClient.for_endpoint(
-            endpoint=endpoint, insecure=insecure, insecure_skip_verify=insecure_skip_verify
-        ),
-        workers=10,
-        max_system_retries=5,
+        client_coro=client_coro,
     )
     return controller