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

flyte/_task_environment.py

@@ -1,30 +1,41 @@
 from __future__ import annotations
 
-import asyncio
+import inspect
 import weakref
 from dataclasses import dataclass, field, replace
 from datetime import timedelta
-from functools import wraps
-from typing import TYPE_CHECKING, Awaitable, Callable, Dict, List, Literal, Optional, ParamSpec, TypeVar, Union, cast
+from typing import (
+    TYPE_CHECKING,
+    Any,
+    Callable,
+    Dict,
+    List,
+    Literal,
+    Optional,
+    Tuple,
+    Union,
+    cast,
+    overload,
+)
 
 import rich.repr
 
-from ._cache import CacheRequest
-from ._datastructures import NativeInterface
+from ._cache import Cache, CacheRequest
 from ._doc import Documentation
 from ._environment import Environment
 from ._image import Image
+from ._link import Link
+from ._pod import PodTemplate
 from ._resources import Resources
 from ._retry import RetryStrategy
 from ._reusable_environment import ReusePolicy
 from ._secret import SecretRequest
 from ._task import AsyncFunctionTaskTemplate, TaskTemplate
+from ._trigger import Trigger
+from .models import MAX_INLINE_IO_BYTES, NativeInterface
 
 if TYPE_CHECKING:
-    from kubernetes.client import V1PodTemplate
-
-P = ParamSpec("P")  # capture the function's parameters
-R = TypeVar("R")  # return type
+    from ._task import F, P, R
 
 
 @rich.repr.auto
@@ -45,99 +56,203 @@ class TaskEnvironment(Environment):
     :param name: Name of the environment
     :param image: Docker image to use for the environment. If set to "auto", will use the default image.
     :param resources: Resources to allocate for the environment.
-    :param env: Environment variables to set for the environment.
+    :param env_vars: Environment variables to set for the environment.
     :param secrets: Secrets to inject into the environment.
-    :param env_dep_hints: Environment dependencies to hint, so when you deploy the environment, the dependencies are
-        also deployed. This is useful when you have a set of environments that depend on each other.
+    :param depends_on: Environment dependencies to hint, so when you deploy the environment,
+        the dependencies are also deployed. This is useful when you have a set of environments
+        that depend on each other.
     :param cache: Cache policy for the environment.
     :param reusable: Reuse policy for the environment, if set, a python process may be reused for multiple tasks.
+    :param plugin_config: Optional plugin configuration for custom task types.
+        If set, all tasks in this environment will use the specified plugin configuration.
+    :param queue: Optional queue name to use for tasks in this environment.
+        If not set, the default queue will be used.
+    :param pod_template: Optional pod template to use for tasks in this environment.
+        If not set, the default pod template will be used.
     """
 
-    cache: Union[CacheRequest] = "auto"
-    reusable: Union[ReusePolicy, Literal["auto"], None] = None
-    # TODO Shall we make this union of string or env? This way we can lookup the env by module/file:name
-    # TODO also we could add list of files that are used by this environment
+    cache: CacheRequest = "disable"
+    reusable: ReusePolicy | None = None
+    plugin_config: Optional[Any] = None
+    queue: Optional[str] = None
 
     _tasks: Dict[str, TaskTemplate] = field(default_factory=dict, init=False)
 
+    def __post_init__(self) -> None:
+        super().__post_init__()
+        if self.reusable is not None and self.plugin_config is not None:
+            raise ValueError("Cannot set plugin_config when environment is reusable.")
+        if self.reusable and not isinstance(self.reusable, ReusePolicy):
+            raise TypeError(f"Expected reusable to be of type ReusePolicy, got {type(self.reusable)}")
+        if self.cache and not isinstance(self.cache, (str, Cache)):
+            raise TypeError(f"Expected cache to be of type str or Cache, got {type(self.cache)}")
+
     def clone_with(
         self,
         name: str,
         image: Optional[Union[str, Image, Literal["auto"]]] = None,
         resources: Optional[Resources] = None,
-        cache: Union[CacheRequest, None] = None,
-        env: Optional[Dict[str, str]] = None,
-        reusable: Union[ReusePolicy, None] = None,
+        env_vars: Optional[Dict[str, str]] = None,
         secrets: Optional[SecretRequest] = None,
-        env_dep_hints: Optional[List[Environment]] = None,
+        depends_on: Optional[List[Environment]] = None,
+        description: Optional[str] = None,
+        interruptible: Optional[bool] = None,
+        **kwargs: Any,
     ) -> TaskEnvironment:
         """
-        Clone the environment with new settings.
+        Clone the TaskEnvironment with new parameters.
+
+        Besides the base environment parameters, you can override kwargs like `cache`, `reusable`, etc.
+
+        :param name: The name of the environment.
+        :param image: The image to use for the environment.
+        :param resources: The resources to allocate for the environment.
+        :param env_vars: The environment variables to set for the environment.
+        :param secrets: The secrets to inject into the environment.
+        :param depends_on: The environment dependencies to hint, so when you deploy the environment,
+            the dependencies are also deployed. This is useful when you have a set of environments
+            that depend on each other.
+        :param queue: The queue name to use for tasks in this environment.
+        :param pod_template: The pod template to use for tasks in this environment.
+        :param description: The description of the environment.
+        :param interruptible: Whether the environment is interruptible and can be scheduled on spot/preemptible
+            instances.
+        :param kwargs: Additional parameters to override the environment (e.g., cache, reusable, plugin_config).
         """
-        if image is None:
-            image = self.image
-        else:
-            image = "auto"
-        return replace(
-            self,
-            cache=cache if cache else "auto",
-            reusable=reusable,
-            name=name,
-            image=image,
-            resources=resources,
-            env=env,
-            secrets=secrets,
-            env_dep_hints=env_dep_hints if env_dep_hints else [],
-        )
-
-    def _task(
+        cache = kwargs.pop("cache", None)
+        reusable = None
+        reusable_set = False
+        if "reusable" in kwargs:
+            reusable_set = True
+            reusable = kwargs.pop("reusable", None)
+
+        # validate unknown kwargs if needed
+        if kwargs:
+            raise TypeError(f"Unexpected keyword arguments: {list(kwargs.keys())}")
+
+        kwargs = self._get_kwargs()
+        kwargs["name"] = name
+        if image is not None:
+            kwargs["image"] = image
+        if resources is not None:
+            kwargs["resources"] = resources
+        if cache is not None:
+            kwargs["cache"] = cache
+        if env_vars is not None:
+            kwargs["env_vars"] = env_vars
+        if reusable_set:
+            kwargs["reusable"] = reusable
+        if secrets is not None:
+            kwargs["secrets"] = secrets
+        if depends_on is not None:
+            kwargs["depends_on"] = depends_on
+        if description is not None:
+            kwargs["description"] = description
+        if interruptible is not None:
+            kwargs["interruptible"] = interruptible
+        return replace(self, **kwargs)
+
+    @overload
+    def task(
         self,
-        _func=None,
         *,
-        name: Optional[str] = None,
-        cache: Union[CacheRequest] | None = None,
+        short_name: Optional[str] = None,
+        cache: CacheRequest | None = None,
         retries: Union[int, RetryStrategy] = 0,
         timeout: Union[timedelta, int] = 0,
         docs: Optional[Documentation] = None,
-        secrets: Optional[SecretRequest] = None,
-        pod_template: Optional[Union[str, "V1PodTemplate"]] = None,
+        pod_template: Optional[Union[str, PodTemplate]] = None,
+        report: bool = False,
+        interruptible: bool | None = None,
+        max_inline_io_bytes: int = MAX_INLINE_IO_BYTES,
+        queue: Optional[str] = None,
+        triggers: Tuple[Trigger, ...] | Trigger = (),
+        links: Tuple[Link, ...] | Link = (),
+    ) -> Callable[[Callable[P, R]], AsyncFunctionTaskTemplate[P, R, Callable[P, R]]]: ...
+
+    @overload
+    def task(
+        self,
+        _func: Callable[P, R],
+        /,
+    ) -> AsyncFunctionTaskTemplate[P, R, Callable[P, R]]: ...
+
+    def task(
+        self,
+        _func: F | None = None,
+        *,
+        short_name: Optional[str] = None,
+        cache: CacheRequest | None = None,
+        retries: Union[int, RetryStrategy] = 0,
+        timeout: Union[timedelta, int] = 0,
+        docs: Optional[Documentation] = None,
+        pod_template: Optional[Union[str, PodTemplate]] = None,
         report: bool = False,
-    ) -> Union[AsyncFunctionTaskTemplate, Callable[P, R]]:
+        interruptible: bool | None = None,
+        max_inline_io_bytes: int = MAX_INLINE_IO_BYTES,
+        queue: Optional[str] = None,
+        triggers: Tuple[Trigger, ...] | Trigger = (),
+        links: Tuple[Link, ...] | Link = (),
+    ) -> Callable[[F], AsyncFunctionTaskTemplate[P, R, F]] | AsyncFunctionTaskTemplate[P, R, F]:
         """
-        :param name: Optional The name of the task (defaults to the function name)
+        Decorate a function to be a task.
+
+        :param _func: Optional The function to decorate. If not provided, the decorator will return a callable that
+        accepts a function to be decorated.
+        :param short_name: Optional A friendly name for the task (defaults to the function name)
         :param cache: Optional The cache policy for the task, defaults to auto, which will cache the results of the
         task.
         :param retries: Optional The number of retries for the task, defaults to 0, which means no retries.
         :param docs: Optional The documentation for the task, if not provided the function docstring will be used.
-        :param secrets: Optional The secrets that will be injected into the task at runtime.
         :param timeout: Optional The timeout for the task.
         :param pod_template: Optional The pod template for the task, if not provided the default pod template will be
         used.
         :param report: Optional Whether to generate the html report for the task, defaults to False.
+        :param max_inline_io_bytes: Maximum allowed size (in bytes) for all inputs and outputs passed directly to the
+         task (e.g., primitives, strings, dicts). Does not apply to files, directories, or dataframes.
+        :param triggers: Optional A tuple of triggers to associate with the task. This allows the task to be run on a
+         schedule or in response to events. Triggers can be defined using the `flyte.trigger` module.
+        :param links: Optional A tuple of links to associate with the task. Links can be used to provide
+         additional context or information about the task. Links should implement the `flyte.Link` protocol
+        :param interruptible: Optional Whether the task is interruptible, defaults to environment setting.
+        :param queue: Optional queue name to use for this task. If not set, the environment's queue will be used.
+
+        :return: A TaskTemplate that can be used to deploy the task.
         """
+        from ._task import F, P, R
+
         if self.reusable is not None:
             if pod_template is not None:
                 raise ValueError("Cannot set pod_template when environment is reusable.")
 
-        def decorator(func: Callable[P, Awaitable[R]]) -> AsyncFunctionTaskTemplate[P, R]:
-            task_name = name or func.__name__
-            task_name = self.name + "." + task_name
-            if len(task_name) > 30:
-                # delete this if we can remove the restriction that task names must be <= 30 characters
-                task_name = task_name[-30:]
-
-            @wraps(func)
-            async def wrapper(*args: P.args, **kwargs: P.kwargs) -> R:
-                return await func(*args, **kwargs)
-
-            if not asyncio.iscoroutinefunction(func):
-                raise TypeError(
-                    f"Function {func.__name__} is not a coroutine function. Use @env.task decorator for async tasks."
-                    f"You can simply mark your function as async def {func.__name__} to make it a coroutine function, "
-                    f"it is ok to write sync code in async functions, but not the other way around."
+        def decorator(func: F) -> AsyncFunctionTaskTemplate[P, R, F]:
+            short = short_name or func.__name__
+            task_name = self.name + "." + func.__name__
+
+            if not inspect.iscoroutinefunction(func) and self.reusable is not None:
+                if self.reusable.concurrency > 1:
+                    raise ValueError(
+                        "Reusable environments with concurrency greater than 1 are only supported for async tasks. "
+                        "Please use an async function or set concurrency to 1."
+                    )
+
+            if self.plugin_config is not None:
+                from flyte.extend import TaskPluginRegistry
+
+                task_template_class: type[AsyncFunctionTaskTemplate[P, R, F]] | None = TaskPluginRegistry.find(
+                    config_type=type(self.plugin_config)
                 )
-            tmpl = AsyncFunctionTaskTemplate(
-                func=wrapper,
+                if task_template_class is None:
+                    raise ValueError(
+                        f"No task plugin found for config type {type(self.plugin_config)}. "
+                        f"Please register a plugin using flyte.extend.TaskPluginRegistry.register() api."
+                    )
+            else:
+                task_template_class = AsyncFunctionTaskTemplate[P, R, F]
+
+            task_template_class = cast(type[AsyncFunctionTaskTemplate[P, R, F]], task_template_class)
+            tmpl = task_template_class(
+                func=func,
                 name=task_name,
                 image=self.image,
                 resources=self.resources,
@@ -146,42 +261,27 @@ class TaskEnvironment(Environment):
                 timeout=timeout,
                 reusable=self.reusable,
                 docs=docs,
-                env=self.env,
-                secrets=secrets or self.secrets,
+                env_vars=self.env_vars,
+                secrets=self.secrets,
                 pod_template=pod_template or self.pod_template,
                 parent_env=weakref.ref(self),
+                parent_env_name=self.name,
                 interface=NativeInterface.from_callable(func),
                 report=report,
+                short_name=short,
+                plugin_config=self.plugin_config,
+                max_inline_io_bytes=max_inline_io_bytes,
+                queue=queue or self.queue,
+                interruptible=interruptible if interruptible is not None else self.interruptible,
+                triggers=triggers if isinstance(triggers, tuple) else (triggers,),
+                links=links if isinstance(links, tuple) else (links,),
             )
             self._tasks[task_name] = tmpl
             return tmpl
 
         if _func is None:
-            return cast(AsyncFunctionTaskTemplate, decorator)
-        return cast(AsyncFunctionTaskTemplate, decorator(_func))
-
-    @property
-    def task(self) -> Callable:
-        """
-        Decorator to create a new task with the environment settings.
-        The task will be executed in its own container with the specified image, resources, and environment variables,
-        unless reusePolicy is set, in which case the same container will be reused for all tasks with the same
-        environment settings.
-
-        :param name: Optional The name of the task (defaults to the function name)
-        :param cache: Optional The cache policy for the task, defaults to auto, which will cache the results of the
-         task.
-        :param retries: Optional The number of retries for the task, defaults to 0, which means no retries.
-        :param docs: Optional The documentation for the task, if not provided the function docstring will be used.
-        :param secrets: Optional The secrets that will be injected into the task at runtime.
-        :param timeout: Optional The timeout for the task.
-        :param pod_template: Optional The pod template for the task, if not provided the default pod template will be
-         used.
-        :param report: Optional Whether to generate the html report for the task, defaults to False.
-
-        :return: New Task instance or Task decorator
-        """
-        return self._task
+            return cast(Callable[[F], AsyncFunctionTaskTemplate[P, R, F]], decorator)
+        return cast(AsyncFunctionTaskTemplate[P, R, F], decorator(_func))
 
     @property
     def tasks(self) -> Dict[str, TaskTemplate]:
@@ -190,11 +290,33 @@ class TaskEnvironment(Environment):
         """
         return self._tasks
 
-    def add_task(self, task: TaskTemplate) -> TaskTemplate:
+    @classmethod
+    def from_task(cls, name: str, *tasks: TaskTemplate) -> TaskEnvironment:
         """
-        Add a task to the environment.
+        Create a TaskEnvironment from a list of tasks. All tasks should have the same image or no Image defined.
+        Similarity of Image is determined by the python reference, not by value.
+
+        If images are different, an error is raised. If no image is defined, the image is set to "auto".
+
+        For any other tasks that need to be use these tasks, the returned environment can be used in the `depends_on`
+        attribute of the other TaskEnvironment.
+
+        :param name: The name of the environment.
+        :param tasks: The list of tasks to create the environment from.
+
+        :raises ValueError: If tasks are assigned to multiple environments or have different images.
+        :return: The created TaskEnvironment.
         """
-        if task.name in self._tasks:
-            raise ValueError(f"Task {task.name} already exists in the environment. Task names should be unique.")
-        self._tasks[task.name] = task
-        return task
+        envs = [t.parent_env() for t in tasks if t.parent_env and t.parent_env() is not None]
+        if envs:
+            raise ValueError("Tasks cannot assigned to multiple environments.")
+        images = {t.image for t in tasks}
+        if len(images) > 1:
+            raise ValueError("Tasks must have the same image to be in the same environment.")
+        image: Union[str, Image] = images.pop() if images else "auto"
+        env = cls(name, image=image)
+        for t in tasks:
+            env._tasks[t.name] = t
+            t.parent_env = weakref.ref(env)
+            t.parent_env_name = name
+        return env

flyte/_task_plugins.py

@@ -0,0 +1,47 @@
+from __future__ import annotations
+
+import typing
+from typing import Type
+
+import rich.repr
+
+if typing.TYPE_CHECKING:
+    from ._task import AsyncFunctionTaskTemplate
+
+T = typing.TypeVar("T", bound="AsyncFunctionTaskTemplate")
+
+
+class _Registry:
+    """
+    A registry for task plugins.
+    """
+
+    def __init__(self: _Registry):
+        self._plugins: typing.Dict[Type, Type[typing.Any]] = {}
+
+    def register(self, config_type: Type, plugin: Type[T]):
+        """
+        Register a plugin.
+        """
+        self._plugins[config_type] = plugin
+
+    def find(self, config_type: Type) -> typing.Optional[Type[T]]:
+        """
+        Get a plugin by name.
+        """
+        return self._plugins.get(config_type)
+
+    def list_plugins(self):
+        """
+        List all registered plugins.
+        """
+        return list(self._plugins.keys())
+
+    def __rich_repr__(self) -> "rich.repr.Result":
+        yield from (("Name", i) for i in self.list_plugins())
+
+    def __repr__(self):
+        return f"TaskPluginRegistry(plugins={self.list_plugins()})"
+
+
+TaskPluginRegistry = _Registry()

flyte/_tools.py

@@ -1,6 +1,3 @@
-import os
-
-
 def ipython_check() -> bool:
     """
     Check if interface is launching from iPython (not colab)
@@ -17,11 +14,14 @@ def ipython_check() -> bool:
     return is_ipython
 
 
-def is_in_cluster() -> bool:
+def ipywidgets_check() -> bool:
     """
-    Check if the task is running in a cluster
-    :return is_in_cluster (bool): True or False
+    Check if the interface is running in IPython with ipywidgets support.
+    :return: True if running in IPython with ipywidgets support, False otherwise.
     """
-    if os.getenv("_UN_CLS"):
+    try:
+        import ipywidgets  # noqa: F401
+
         return True
-    return False
+    except (ImportError, NameError):
+        return False

flyte/_trace.py

@@ -1,10 +1,10 @@
 import functools
 import inspect
 import time
-from datetime import timedelta
 from typing import Any, AsyncGenerator, AsyncIterator, Awaitable, Callable, TypeGuard, TypeVar, Union, cast
 
-from flyte._datastructures import NativeInterface
+from flyte._logging import logger
+from flyte.models import NativeInterface
 
 T = TypeVar("T")
 
@@ -14,7 +14,6 @@ def trace(func: Callable[..., T]) -> Callable[..., T]:
     A decorator that traces function execution with timing information.
     Works with regular functions, async functions, and async generators/iterators.
     """
-    func_name = func.__name__
 
     @functools.wraps(func)
     def wrapper_sync(*args: Any, **kwargs: Any) -> Any:
@@ -31,27 +30,29 @@ def trace(func: Callable[..., T]) -> Callable[..., T]:
             # We will also check if we are not initialized, It is not expected to be not initialized
             from ._internal.controllers import get_controller
 
-            controller = await get_controller()
+            controller = get_controller()
             iface = NativeInterface.from_callable(func)
-            info, ok = await controller.get_action_outputs(iface, func_name, *args, **kwargs)
+            info, ok = await controller.get_action_outputs(iface, func, *args, **kwargs)
             if ok:
+                logger.info(f"Found existing trace info for {func}, {info}")
                 if info.output:
                     return info.output
                 elif info.error:
                     raise info.error
+            else:
+                logger.debug(f"No existing trace info found for {func}, proceeding to execute.")
             start_time = time.time()
             try:
                 # Cast to Awaitable to satisfy mypy
                 coroutine_result = cast(Awaitable[Any], func(*args, **kwargs))
                 results = await coroutine_result
-                duration = time.time() - start_time
-                info.add_outputs(results, timedelta(seconds=duration))
+                info.add_outputs(results, start_time=start_time, end_time=time.time())
                 await controller.record_trace(info)
+                logger.debug(f"Finished trace for {func}, {info}")
                 return results
             except Exception as e:
                 # If there is an error, we need to record it
-                duration = time.time() - start_time
-                info.add_error(e, timedelta(seconds=duration))
+                info.add_error(e, start_time=start_time, end_time=time.time())
                 await controller.record_trace(info)
                 raise e
         else:
@@ -74,9 +75,9 @@ def trace(func: Callable[..., T]) -> Callable[..., T]:
             # We will also check if we are not initialized, It is not expected to be not initialized
             from ._internal.controllers import get_controller
 
-            controller = await get_controller()
+            controller = get_controller()
             iface = NativeInterface.from_callable(func)
-            info, ok = await controller.get_action_outputs(iface, func_name, *args, **kwargs)
+            info, ok = await controller.get_action_outputs(iface, func, *args, **kwargs)
             if ok:
                 if info.output:
                     for item in info.output:
@@ -88,25 +89,17 @@ def trace(func: Callable[..., T]) -> Callable[..., T]:
                 items = []
                 result = func(*args, **kwargs)
                 # TODO ideally we should use streaming into the type-engine so that it stream uploads large blocks
-                if inspect.isasyncgen(result):
+                if inspect.isasyncgen(result) or is_async_iterable(result):
                     # If it's directly an async generator
                     async_iter = result
                     async for item in async_iter:
                         items.append(item)
                         yield item
-                elif is_async_iterable(result):
-                    # If it's an async iterable (has __aiter__)
-                    async for item in result:
-                        items.append(item)
-                        yield item
-                duration = time.time() - start_time
-                info.add_outputs(items, timedelta(seconds=duration))
+                info.add_outputs(items, start_time=start_time, end_time=time.time())
                 await controller.record_trace(info)
                 return
             except Exception as e:
-                end_time = time.time()
-                duration = end_time - start_time
-                info.add_error(e, timedelta(seconds=duration))
+                info.add_error(e, start_time=start_time, end_time=time.time())
                 await controller.record_trace(info)
                 raise e
         else:
@@ -118,10 +111,8 @@ def trace(func: Callable[..., T]) -> Callable[..., T]:
     # Choose the appropriate wrapper based on the function type
     if inspect.iscoroutinefunction(func):
         # This handles async functions that return normal values
-        print(f"Coroutine function {func.__name__}")
         return cast(Callable[..., T], wrapper_async)
     elif inspect.isasyncgenfunction(func):
-        print(f"Async generator function {func.__name__}")
         return cast(Callable[..., T], wrapper_async_iterator)
     else:
         # For regular sync functions