flyte 2.0.0b13__py3-none-any.whl → 2.0.0b30__py3-none-any.whl

flyte/_task_environment.py

@@ -12,27 +12,29 @@ from typing import (
     List,
     Literal,
     Optional,
+    Tuple,
     Union,
     cast,
+    overload,
 )
 
 import rich.repr
 
-from ._cache import CacheRequest
+from ._cache import Cache, CacheRequest
 from ._doc import Documentation
 from ._environment import Environment
 from ._image import Image
+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
-
-    from ._task import FunctionTypes, P, R
+    from ._task import F, P, R
 
 
 @rich.repr.auto
@@ -60,13 +62,18 @@ class TaskEnvironment(Environment):
         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: CacheRequest = "disable"
     reusable: ReusePolicy | None = None
     plugin_config: Optional[Any] = 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
+    queue: Optional[str] = None
 
     _tasks: Dict[str, TaskTemplate] = field(default_factory=dict, init=False)
 
@@ -74,6 +81,10 @@ class TaskEnvironment(Environment):
         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,
@@ -83,6 +94,8 @@ class TaskEnvironment(Environment):
         env_vars: Optional[Dict[str, str]] = None,
         secrets: Optional[SecretRequest] = None,
         depends_on: Optional[List[Environment]] = None,
+        description: Optional[str] = None,
+        interruptible: Optional[bool] = None,
         **kwargs: Any,
     ) -> TaskEnvironment:
         """
@@ -98,6 +111,11 @@ class TaskEnvironment(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).
         """
         cache = kwargs.pop("cache", None)
@@ -127,27 +145,58 @@ class TaskEnvironment(Environment):
             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,
+        *,
+        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,
+        interruptible: bool | None = None,
+        max_inline_io_bytes: int = MAX_INLINE_IO_BYTES,
+        queue: Optional[str] = None,
+        triggers: Tuple[Trigger, ...] | Trigger = (),
+    ) -> 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=None,
+        _func: F | None = None,
         *,
-        name: Optional[str] = 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, "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,
-    ) -> Union[AsyncFunctionTaskTemplate, Callable[P, R]]:
+        queue: Optional[str] = None,
+        triggers: Tuple[Trigger, ...] | Trigger = (),
+    ) -> Callable[[F], AsyncFunctionTaskTemplate[P, R, F]] | AsyncFunctionTaskTemplate[P, R, F]:
         """
         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 name: Optional A friendly name for the task (defaults to the function name)
+        :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.
@@ -158,15 +207,21 @@ class TaskEnvironment(Environment):
         :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 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 P, R
+        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: FunctionTypes) -> AsyncFunctionTaskTemplate[P, R]:
-            friendly_name = name or func.__name__
+        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:
@@ -179,7 +234,7 @@ class TaskEnvironment(Environment):
             if self.plugin_config is not None:
                 from flyte.extend import TaskPluginRegistry
 
-                task_template_class: type[AsyncFunctionTaskTemplate[P, R]] | None = TaskPluginRegistry.find(
+                task_template_class: type[AsyncFunctionTaskTemplate[P, R, F]] | None = TaskPluginRegistry.find(
                     config_type=type(self.plugin_config)
                 )
                 if task_template_class is None:
@@ -188,9 +243,9 @@ class TaskEnvironment(Environment):
                         f"Please register a plugin using flyte.extend.TaskPluginRegistry.register() api."
                     )
             else:
-                task_template_class = AsyncFunctionTaskTemplate[P, R]
+                task_template_class = AsyncFunctionTaskTemplate[P, R, F]
 
-            task_template_class = cast(type[AsyncFunctionTaskTemplate[P, R]], task_template_class)
+            task_template_class = cast(type[AsyncFunctionTaskTemplate[P, R, F]], task_template_class)
             tmpl = task_template_class(
                 func=func,
                 name=task_name,
@@ -205,18 +260,22 @@ class TaskEnvironment(Environment):
                 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,
-                friendly_name=friendly_name,
+                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,),
             )
             self._tasks[task_name] = tmpl
             return tmpl
 
         if _func is None:
-            return cast(AsyncFunctionTaskTemplate, decorator)
-        return cast(AsyncFunctionTaskTemplate, decorator(_func))
+            return cast(Callable[[F], AsyncFunctionTaskTemplate[P, R, F]], decorator)
+        return cast(AsyncFunctionTaskTemplate[P, R, F], decorator(_func))
 
     @property
     def tasks(self) -> Dict[str, TaskTemplate]:
@@ -225,16 +284,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".
 
-        Useful when you want to add a task to an environment that is not defined using the `task` decorator.
+        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.
 
-        :param task: The TaskTemplate to add to this environment.
+        :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
-        task.parent_env = weakref.ref(self)
-        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

@@ -1,3 +1,5 @@
+from __future__ import annotations
+
 import typing
 from typing import Type
 
@@ -14,8 +16,8 @@ class _Registry:
     A registry for task plugins.
     """
 
-    def __init__(self):
-        self._plugins: typing.Dict[Type, Type[T]] = {}
+    def __init__(self: _Registry):
+        self._plugins: typing.Dict[Type, Type[typing.Any]] = {}
 
     def register(self, config_type: Type, plugin: Type[T]):
         """

flyte/_trace.py

@@ -3,6 +3,7 @@ import inspect
 import time
 from typing import Any, AsyncGenerator, AsyncIterator, Awaitable, Callable, TypeGuard, TypeVar, Union, cast
 
+from flyte._logging import logger
 from flyte.models import NativeInterface
 
 T = TypeVar("T")
@@ -33,10 +34,13 @@ def trace(func: Callable[..., T]) -> Callable[..., T]:
             iface = NativeInterface.from_callable(func)
             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
@@ -44,6 +48,7 @@ def trace(func: Callable[..., T]) -> Callable[..., T]:
                 results = await coroutine_result
                 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