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

flyte/_task.py

@@ -15,10 +15,12 @@ from typing import (
     Literal,
     Optional,
     ParamSpec,
+    Tuple,
     TypeAlias,
     TypeVar,
     Union,
     cast,
+    overload,
 )
 
 from flyte._pod import PodTemplate
@@ -33,10 +35,11 @@ from ._retry import RetryStrategy
 from ._reusable_environment import ReusePolicy
 from ._secret import SecretRequest
 from ._timeout import TimeoutType
+from ._trigger import Trigger
 from .models import MAX_INLINE_IO_BYTES, NativeInterface, SerializationContext
 
 if TYPE_CHECKING:
-    from flyteidl.core.tasks_pb2 import DataLoadingConfig
+    from flyteidl2.core.tasks_pb2 import DataLoadingConfig
 
     from ._task_environment import TaskEnvironment
 
@@ -45,11 +48,12 @@ R = TypeVar("R")  # return type
 
 AsyncFunctionType: TypeAlias = Callable[P, Coroutine[Any, Any, R]]
 SyncFunctionType: TypeAlias = Callable[P, R]
-FunctionTypes: TypeAlias = Union[AsyncFunctionType, SyncFunctionType]
+FunctionTypes: TypeAlias = AsyncFunctionType | SyncFunctionType
+F = TypeVar("F", bound=FunctionTypes)
 
 
 @dataclass(kw_only=True)
-class TaskTemplate(Generic[P, R]):
+class TaskTemplate(Generic[P, R, F]):
     """
     Task template is a template for a task that can be executed. It defines various parameters for the task, which
     can be defined statically at the time of task definition or dynamically at the time of task invocation using
@@ -69,8 +73,8 @@ class TaskTemplate(Generic[P, R]):
     version with flyte installed
     :param resources: Optional The resources to use for the task
     :param cache: Optional The cache policy for the task, defaults to auto, which will cache the results of the task.
-    :param interruptable: Optional The interruptable policy for the task, defaults to False, which means the task
-     will not be scheduled on interruptable nodes. If set to True, the task will be scheduled on interruptable nodes,
+    :param interruptible: Optional The interruptible policy for the task, defaults to False, which means the task
+     will not be scheduled on interruptible nodes. If set to True, the task will be scheduled on interruptible nodes,
      and the code should handle interruptions and resumptions.
     :param retries: Optional The number of retries for the task, defaults to 0, which means no retries.
     :param reusable: Optional The reusability policy for the task, defaults to None, which means the task environment
@@ -81,6 +85,10 @@ class TaskTemplate(Generic[P, R]):
     :param timeout: Optional The timeout for the task.
     :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 pod_template: Optional The pod template to use for the task.
+    :param report: Optional Whether to report the task execution to the Flyte console, defaults to False.
+    :param queue: Optional The queue to use for the task. If not provided, the default queue will be used.
+    :param debuggable: Optional Whether the task supports debugging capabilities, defaults to False.
     """
 
     name: str
@@ -90,8 +98,8 @@ class TaskTemplate(Generic[P, R]):
     task_type_version: int = 0
     image: Union[str, Image, Literal["auto"]] = "auto"
     resources: Optional[Resources] = None
-    cache: CacheRequest = "auto"
-    interruptable: bool = False
+    cache: CacheRequest = "disable"
+    interruptible: bool = False
     retries: Union[int, RetryStrategy] = 0
     reusable: Union[ReusePolicy, None] = None
     docs: Optional[Documentation] = None
@@ -100,10 +108,14 @@ class TaskTemplate(Generic[P, R]):
     timeout: Optional[TimeoutType] = None
     pod_template: Optional[Union[str, PodTemplate]] = None
     report: bool = False
+    queue: Optional[str] = None
+    debuggable: bool = False
 
     parent_env: Optional[weakref.ReferenceType[TaskEnvironment]] = None
+    parent_env_name: Optional[str] = None
     ref: bool = field(default=False, init=False, repr=False, compare=False)
     max_inline_io_bytes: int = MAX_INLINE_IO_BYTES
+    triggers: Tuple[Trigger, ...] = field(default_factory=tuple)
 
     # Only used in python 3.10 and 3.11, where we cannot use markcoroutinefunction
     _call_as_synchronous: bool = False
@@ -217,6 +229,14 @@ class TaskTemplate(Generic[P, R]):
     def native_interface(self) -> NativeInterface:
         return self.interface
 
+    @overload
+    async def aio(self: TaskTemplate[P, R, SyncFunctionType], *args: P.args, **kwargs: P.kwargs) -> R: ...
+
+    @overload
+    async def aio(
+        self: TaskTemplate[P, R, AsyncFunctionType], *args: P.args, **kwargs: P.kwargs
+    ) -> Coroutine[Any, Any, R]: ...
+
     async def aio(self, *args: P.args, **kwargs: P.kwargs) -> Coroutine[Any, Any, R] | R:
         """
         The aio function allows executing "sync" tasks, in an async context. This helps with migrating v1 defined sync
@@ -240,7 +260,6 @@ class TaskTemplate(Generic[P, R]):
         :param kwargs:
         :return:
         """
-
         ctx = internal_ctx()
         if ctx.is_task_context():
             from ._internal.controllers import get_controller
@@ -265,6 +284,14 @@ class TaskTemplate(Generic[P, R]):
             # even for synchronous tasks. This is to support migration.
             return self.forward(*args, **kwargs)
 
+    @overload
+    def __call__(self: TaskTemplate[P, R, SyncFunctionType], *args: P.args, **kwargs: P.kwargs) -> R: ...
+
+    @overload
+    def __call__(
+        self: TaskTemplate[P, R, AsyncFunctionType], *args: P.args, **kwargs: P.kwargs
+    ) -> Coroutine[Any, Any, R]: ...
+
     def __call__(self, *args: P.args, **kwargs: P.kwargs) -> Coroutine[Any, Any, R] | R:
         """
         This is the entrypoint for an async function task at runtime. It will be called during an execution.
@@ -327,11 +354,30 @@ class TaskTemplate(Generic[P, R]):
         secrets: Optional[SecretRequest] = None,
         max_inline_io_bytes: int | None = None,
         pod_template: Optional[Union[str, PodTemplate]] = None,
+        queue: Optional[str] = None,
+        interruptible: Optional[bool] = None,
         **kwargs: Any,
     ) -> TaskTemplate:
         """
         Override various parameters of the task template. This allows for dynamic configuration of the task
         when it is called, such as changing the image, resources, cache policy, etc.
+
+        :param short_name: Optional override for the short name of the task.
+        :param resources: Optional override for the resources to use for the task.
+        :param cache: Optional override for the cache policy for the task.
+        :param retries: Optional override for the number of retries for the task.
+        :param timeout: Optional override for the timeout for the task.
+        :param reusable: Optional override for the reusability policy for the task.
+        :param env_vars: Optional override for the environment variables to set for the task.
+        :param secrets: Optional override for the secrets that will be injected into the task at runtime.
+        :param max_inline_io_bytes: Optional override for the maximum allowed size (in bytes) for all inputs and outputs
+         passed directly to the task.
+        :param pod_template: Optional override for the pod template to use for the task.
+        :param queue: Optional override for the queue to use for the task.
+        :param kwargs: Additional keyword arguments for further overrides. Some fields like name, image, docs,
+         and interface cannot be overridden.
+
+        :return: A new TaskTemplate instance with the overridden parameters.
         """
         cache = cache or self.cache
         retries = retries or self.retries
@@ -366,6 +412,8 @@ class TaskTemplate(Generic[P, R]):
         env_vars = env_vars or self.env_vars
         secrets = secrets or self.secrets
 
+        interruptible = interruptible if interruptible is not None else self.interruptible
+
         for k, v in kwargs.items():
             if k == "name":
                 raise ValueError("Name cannot be overridden")
@@ -388,19 +436,22 @@ class TaskTemplate(Generic[P, R]):
             secrets=secrets,
             max_inline_io_bytes=max_inline_io_bytes,
             pod_template=pod_template,
+            interruptible=interruptible,
+            queue=queue or self.queue,
             **kwargs,
         )
 
 
 @dataclass(kw_only=True)
-class AsyncFunctionTaskTemplate(TaskTemplate[P, R]):
+class AsyncFunctionTaskTemplate(TaskTemplate[P, R, F]):
     """
     A task template that wraps an asynchronous functions. This is automatically created when an asynchronous function
     is decorated with the task decorator.
     """
 
-    func: FunctionTypes
+    func: F
     plugin_config: Optional[Any] = None  # This is used to pass plugin specific configuration
+    debuggable: bool = True
 
     def __post_init__(self):
         super().__post_init__()
@@ -479,6 +530,11 @@ class AsyncFunctionTaskTemplate(TaskTemplate[P, R]):
 
             from flyte._internal.resolvers.default import DefaultTaskResolver
 
+            if not serialize_context.root_dir:
+                raise RuntimeSystemError(
+                    "SerializationError",
+                    "Root dir is required for default task resolver when no code bundle is provided.",
+                )
             _task_resolver = DefaultTaskResolver()
             args = [
                 *args,

flyte/_task_environment.py

@@ -12,8 +12,10 @@ from typing import (
     List,
     Literal,
     Optional,
+    Tuple,
     Union,
     cast,
+    overload,
 )
 
 import rich.repr
@@ -22,17 +24,17 @@ 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)
 
@@ -87,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:
         """
@@ -102,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)
@@ -131,21 +145,52 @@ 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,
-        _func: Callable[P, R] | 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, "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[[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,
+        interruptible: bool | None = None,
+        max_inline_io_bytes: int = MAX_INLINE_IO_BYTES,
+        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.
 
@@ -162,14 +207,20 @@ 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]:
+        def decorator(func: F) -> AsyncFunctionTaskTemplate[P, R, F]:
             short = short_name or func.__name__
             task_name = self.name + "." + func.__name__
 
@@ -183,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:
@@ -192,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,
@@ -209,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,
                 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]:
@@ -229,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]):
         """