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

flyte/_task_environment.py

@@ -12,6 +12,7 @@ from typing import (
     List,
     Literal,
     Optional,
+    Tuple,
     Union,
     cast,
 )
@@ -22,16 +23,16 @@ 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
 
 
@@ -60,13 +61,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 +93,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 +110,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,6 +144,10 @@ 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)
 
     def task(
@@ -142,10 +159,13 @@ class TaskEnvironment(Environment):
         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 = (),
+    ) -> Union[AsyncFunctionTaskTemplate[P, R], Callable[P, R]]:
         """
         Decorate a function to be a task.
 
@@ -162,6 +182,12 @@ 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
 
@@ -214,6 +240,9 @@ class TaskEnvironment(Environment):
                 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
@@ -229,16 +258,32 @@ 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)
+        return env

flyte/_trigger.py

@@ -0,0 +1,382 @@
+from __future__ import annotations
+
+from dataclasses import dataclass
+from datetime import datetime
+from typing import Any, Dict, Mapping, Union
+
+import rich.repr
+
+
+class _trigger_time:
+    """
+    This class represents the actual time of Trigger, which can be bound to any task input.
+    """
+
+
+TriggerTime = _trigger_time()
+
+
+@rich.repr.auto
+@dataclass(frozen=True)
+class Cron:
+    """
+    This class defines a Cron automation that can be associated with a Trigger in Flyte.
+    Example usage:
+    ```python
+    from flyte.trigger import Trigger, Cron
+    my_trigger = Trigger(
+        name="my_cron_trigger",
+        automation=Cron("0 * * * *"),  # Runs every hour
+        description="A trigger that runs every hour",
+    )
+    ```
+    """
+
+    expression: str
+
+    def __str__(self):
+        return f"Cron Trigger: {self.expression}"
+
+
+@rich.repr.auto
+@dataclass(frozen=True)
+class FixedRate:
+    """
+    This class defines a FixedRate automation that can be associated with a Trigger in Flyte.
+    Example usage:
+    ```python
+    from flyte.trigger import Trigger, FixedRate
+    my_trigger = Trigger(
+        name="my_fixed_rate_trigger",
+        automation=FixedRate(60),  # Runs every hour
+        description="A trigger that runs every hour",
+    )
+    ```
+    """
+
+    interval_minutes: int
+    start_time: datetime | None = None
+
+    def __str__(self):
+        return f"FixedRate Trigger: every {self.interval_minutes} minutes"
+
+
+@rich.repr.auto
+@dataclass(frozen=True)
+class Trigger:
+    """
+    This class defines specification of a Trigger, that can be associated with any Flyte V2 task.
+    The trigger then is deployed to the Flyte Platform.
+
+    Triggers can be used to run tasks on a schedule, in response to events, or based on other conditions.
+    The `Trigger` class encapsulates the metadata and configuration needed to define a trigger.
+
+    You can associate the same Trigger object with multiple tasks.
+
+    Example usage:
+    ```python
+    from flyte.trigger import Trigger
+    my_trigger = Trigger(
+        name="my_trigger",
+        description="A trigger that runs every hour",
+    )
+    ```
+
+    :param name: (str) The name of the trigger.
+    :param automation: (AutomationType) The automation type, currently only supports Cron.
+    :param description: (str) A description of the trigger, default is an empty string.
+    :param auto_activate: (bool) Whether the trigger should be automatically activated, default is True.
+    :param inputs: (Dict[str, Any]) Optional inputs for the trigger, default is None. If provided, will replace the
+       values for inputs to these defaults.
+    :param env_vars: (Dict[str, str]) Optional environment variables for the trigger, default is None. If provided, will
+        replace the environment variables set in the config of the task.
+    :param interruptible: (bool) Whether the trigger is interruptible, default is None. If provided,
+     it overrides whatever is set in the config of the task.
+    :param overwrite_cache: (bool) Whether to overwrite the cache, default is False.
+    :param queue: (str) Optional queue to run the trigger in, default is None.
+    :param labels: (Mapping[str, str]) Optional labels to attach to the trigger, default is None.
+    :param annotations: (Mapping[str, str]) Optional annotations to attach to the trigger, default is None.
+    """
+
+    name: str
+    automation: Union[Cron, FixedRate]
+    description: str = ""
+    auto_activate: bool = True
+    inputs: Dict[str, Any] | None = None
+    env_vars: Dict[str, str] | None = None
+    interruptible: bool | None = None
+    overwrite_cache: bool = False
+    queue: str | None = None
+    labels: Mapping[str, str] | None = None
+    annotations: Mapping[str, str] | None = None
+
+    def __post_init__(self):
+        if not self.name:
+            raise ValueError("Trigger name cannot be empty")
+        if self.automation is None:
+            raise ValueError("Automation cannot be None")
+
+    @classmethod
+    def daily(
+        cls,
+        trigger_time_input_key: str = "trigger_time",
+        *,
+        name: str = "daily",
+        description: str = "A trigger that runs daily at midnight",
+        auto_activate: bool = True,
+        inputs: Dict[str, Any] | None = None,
+        env_vars: Dict[str, str] | None = None,
+        interruptible: bool | None = None,
+        overwrite_cache: bool = False,
+        queue: str | None = None,
+        labels: Mapping[str, str] | None = None,
+        annotations: Mapping[str, str] | None = None,
+    ) -> Trigger:
+        """
+        Creates a Cron trigger that runs daily at midnight.
+
+        Args:
+            trigger_time_input_key (str): The input key for the trigger time, default is "trigger_time".
+            name (str): The name of the trigger, default is "daily".
+            description (str): A description of the trigger.
+            auto_activate (bool): Whether the trigger should be automatically activated.
+            inputs (Dict[str, Any] | None): Optional inputs for the trigger.
+            env_vars (Dict[str, str] | None): Optional environment variables.
+            interruptible (bool | None): Whether the trigger is interruptible.
+            overwrite_cache (bool): Whether to overwrite the cache.
+            queue (str | None): Optional queue to run the trigger in.
+            labels (Mapping[str, str] | None): Optional labels to attach to the trigger.
+            annotations (Mapping[str, str] | None): Optional annotations to attach to the trigger.
+
+        Returns:
+            Trigger: A trigger that runs daily at midnight.
+        """
+        final_inputs = {trigger_time_input_key: TriggerTime}
+        if inputs:
+            final_inputs.update(inputs)
+
+        return cls(
+            name=name,
+            automation=Cron("0 0 * * *"),  # Cron expression for daily at midnight
+            description=description,
+            auto_activate=auto_activate,
+            inputs=final_inputs,
+            env_vars=env_vars,
+            interruptible=interruptible,
+            overwrite_cache=overwrite_cache,
+            queue=queue,
+            labels=labels,
+            annotations=annotations,
+        )
+
+    @classmethod
+    def hourly(
+        cls,
+        trigger_time_input_key: str = "trigger_time",
+        *,
+        name: str = "hourly",
+        description: str = "A trigger that runs every hour",
+        auto_activate: bool = True,
+        inputs: Dict[str, Any] | None = None,
+        env_vars: Dict[str, str] | None = None,
+        interruptible: bool | None = None,
+        overwrite_cache: bool = False,
+        queue: str | None = None,
+        labels: Mapping[str, str] | None = None,
+        annotations: Mapping[str, str] | None = None,
+    ) -> Trigger:
+        """
+        Creates a Cron trigger that runs every hour.
+
+        Args:
+            trigger_time_input_key (str): The input parameter for the trigger time, default is "trigger_time".
+            name (str): The name of the trigger, default is "hourly".
+            description (str): A description of the trigger.
+            auto_activate (bool): Whether the trigger should be automatically activated.
+            inputs (Dict[str, Any] | None): Optional inputs for the trigger.
+            env_vars (Dict[str, str] | None): Optional environment variables.
+            interruptible (bool | None): Whether the trigger is interruptible.
+            overwrite_cache (bool): Whether to overwrite the cache.
+            queue (str | None): Optional queue to run the trigger in.
+            labels (Mapping[str, str] | None): Optional labels to attach to the trigger.
+            annotations (Mapping[str, str] | None): Optional annotations to attach to the trigger.
+
+        Returns:
+            Trigger: A trigger that runs every hour, on the hour.
+        """
+        final_inputs = {trigger_time_input_key: TriggerTime}
+        if inputs:
+            final_inputs.update(inputs)
+
+        return cls(
+            name=name,
+            automation=Cron("0 * * * *"),  # Cron expression for every hour
+            description=description,
+            auto_activate=auto_activate,
+            inputs=final_inputs,
+            env_vars=env_vars,
+            interruptible=interruptible,
+            overwrite_cache=overwrite_cache,
+            queue=queue,
+            labels=labels,
+            annotations=annotations,
+        )
+
+    @classmethod
+    def minutely(
+        cls,
+        trigger_time_input_key: str = "trigger_time",
+        *,
+        name: str = "minutely",
+        description: str = "A trigger that runs every minute",
+        auto_activate: bool = True,
+        inputs: Dict[str, Any] | None = None,
+        env_vars: Dict[str, str] | None = None,
+        interruptible: bool | None = None,
+        overwrite_cache: bool = False,
+        queue: str | None = None,
+        labels: Mapping[str, str] | None = None,
+        annotations: Mapping[str, str] | None = None,
+    ) -> Trigger:
+        """
+        Creates a Cron trigger that runs every minute.
+
+        Args:
+            trigger_time_input_key (str): The input parameter for the trigger time, default is "trigger_time".
+            name (str): The name of the trigger, default is "every_minute".
+            description (str): A description of the trigger.
+            auto_activate (bool): Whether the trigger should be automatically activated.
+            inputs (Dict[str, Any] | None): Optional inputs for the trigger.
+            env_vars (Dict[str, str] | None): Optional environment variables.
+            interruptible (bool | None): Whether the trigger is interruptible.
+            overwrite_cache (bool): Whether to overwrite the cache.
+            queue (str | None): Optional queue to run the trigger in.
+            labels (Mapping[str, str] | None): Optional labels to attach to the trigger.
+            annotations (Mapping[str, str] | None): Optional annotations to attach to the trigger.
+
+        Returns:
+            Trigger: A trigger that runs every minute.
+        """
+        final_inputs = {trigger_time_input_key: TriggerTime}
+        if inputs:
+            final_inputs.update(inputs)
+
+        return cls(
+            name=name,
+            automation=Cron("* * * * *"),  # Cron expression for every minute
+            description=description,
+            auto_activate=auto_activate,
+            inputs=final_inputs,
+            env_vars=env_vars,
+            interruptible=interruptible,
+            overwrite_cache=overwrite_cache,
+            queue=queue,
+            labels=labels,
+            annotations=annotations,
+        )
+
+    @classmethod
+    def weekly(
+        cls,
+        trigger_time_input_key: str = "trigger_time",
+        *,
+        name: str = "weekly",
+        description: str = "A trigger that runs weekly on Sundays at midnight",
+        auto_activate: bool = True,
+        inputs: Dict[str, Any] | None = None,
+        env_vars: Dict[str, str] | None = None,
+        interruptible: bool | None = None,
+        overwrite_cache: bool = False,
+        queue: str | None = None,
+        labels: Mapping[str, str] | None = None,
+        annotations: Mapping[str, str] | None = None,
+    ) -> Trigger:
+        """
+        Creates a Cron trigger that runs weekly on Sundays at midnight.
+
+        Args:
+            trigger_time_input_key (str): The input parameter for the trigger time, default is "trigger_time".
+            name (str): The name of the trigger, default is "weekly".
+            description (str): A description of the trigger.
+            auto_activate (bool): Whether the trigger should be automatically activated.
+            inputs (Dict[str, Any] | None): Optional inputs for the trigger.
+            env_vars (Dict[str, str] | None): Optional environment variables.
+            interruptible (bool | None): Whether the trigger is interruptible.
+            overwrite_cache (bool): Whether to overwrite the cache.
+            queue (str | None): Optional queue to run the trigger in.
+            labels (Mapping[str, str] | None): Optional labels to attach to the trigger.
+            annotations (Mapping[str, str] | None): Optional annotations to attach to the trigger.
+
+        Returns:
+            Trigger: A trigger that runs weekly on Sundays at midnight.
+        """
+        final_inputs = {trigger_time_input_key: TriggerTime}
+        if inputs:
+            final_inputs.update(inputs)
+
+        return cls(
+            name=name,
+            automation=Cron("0 0 * * 0"),  # Cron expression for every Sunday at midnight
+            description=description,
+            auto_activate=auto_activate,
+            inputs=final_inputs,
+            env_vars=env_vars,
+            interruptible=interruptible,
+            overwrite_cache=overwrite_cache,
+            queue=queue,
+            labels=labels,
+            annotations=annotations,
+        )
+
+    @classmethod
+    def monthly(
+        cls,
+        trigger_time_input_key: str = "trigger_time",
+        *,
+        name: str = "monthly",
+        description: str = "A trigger that runs monthly on the 1st at midnight",
+        auto_activate: bool = True,
+        inputs: Dict[str, Any] | None = None,
+        env_vars: Dict[str, str] | None = None,
+        interruptible: bool | None = None,
+        overwrite_cache: bool = False,
+        queue: str | None = None,
+        labels: Mapping[str, str] | None = None,
+        annotations: Mapping[str, str] | None = None,
+    ) -> Trigger:
+        """
+        Creates a Cron trigger that runs monthly on the 1st at midnight.
+
+        Args:
+            trigger_time_input_key (str): The input parameter for the trigger time, default is "trigger_time".
+            name (str): The name of the trigger, default is "monthly".
+            description (str): A description of the trigger.
+            auto_activate (bool): Whether the trigger should be automatically activated.
+            inputs (Dict[str, Any] | None): Optional inputs for the trigger.
+            env_vars (Dict[str, str] | None): Optional environment variables.
+            interruptible (bool | None): Whether the trigger is interruptible.
+            overwrite_cache (bool): Whether to overwrite the cache.
+            queue (str | None): Optional queue to run the trigger in.
+            labels (Mapping[str, str] | None): Optional labels to attach to the trigger.
+            annotations (Mapping[str, str] | None): Optional annotations to attach to the trigger.
+
+        Returns:
+            Trigger: A trigger that runs monthly on the 1st at midnight.
+        """
+        final_inputs = {trigger_time_input_key: TriggerTime}
+        if inputs:
+            final_inputs.update(inputs)
+
+        return cls(
+            name=name,
+            automation=Cron("0 0 1 * *"),  # Cron expression for monthly on the 1st at midnight
+            description=description,
+            auto_activate=auto_activate,
+            inputs=final_inputs,
+            env_vars=env_vars,
+            interruptible=interruptible,
+            overwrite_cache=overwrite_cache,
+            queue=queue,
+            labels=labels,
+            annotations=annotations,
+        )

flyte/_version.py

@@ -28,7 +28,7 @@ version_tuple: VERSION_TUPLE
 commit_id: COMMIT_ID
 __commit_id__: COMMIT_ID
 
-__version__ = version = '2.0.0b22'
-__version_tuple__ = version_tuple = (2, 0, 0, 'b22')
+__version__ = version = '2.0.0b23'
+__version_tuple__ = version_tuple = (2, 0, 0, 'b23')
 
-__commit_id__ = commit_id = 'gce9a6bede'
+__commit_id__ = commit_id = 'g52dc01de2'

flyte/cli/_abort.py

@@ -1,5 +1,4 @@
 import rich_click as click
-from rich.console import Console
 
 from flyte.cli import _common as common
 
@@ -23,6 +22,7 @@ def run(cfg: common.CLIConfig, run_name: str, project: str | None = None, domain
     cfg.init(project=project, domain=domain)
     r = Run.get(name=run_name)
     if r:
-        console = Console()
-        r.abort()
+        console = common.get_console()
+        with console.status(f"Aborting run '{run_name}'...", spinner="dots"):
+            r.abort()
         console.print(f"Run '{run_name}' has been aborted.")

flyte/cli/_build.py

@@ -44,9 +44,7 @@ class BuildEnvCommand(click.Command):
         super().__init__(*args, **kwargs)
 
     def invoke(self, ctx: click.Context):
-        from rich.console import Console
-
-        console = Console()
+        console = common.get_console()
         console.print(f"Building Environment: {self.obj_name}")
         obj: CLIConfig = ctx.obj
         obj.init()

flyte/cli/_common.py

@@ -1,6 +1,7 @@
 from __future__ import annotations
 
 import importlib.util
+import json
 import logging
 import os
 import sys
@@ -22,7 +23,7 @@ from rich.traceback import Traceback
 import flyte.errors
 from flyte.config import Config
 
-OutputFormat = Literal["table", "json", "table-simple"]
+OutputFormat = Literal["table", "json", "table-simple", "json-raw"]
 
 PREFERRED_BORDER_COLOR = "dim cyan"
 PREFERRED_ACCENT_COLOR = "bold #FFD700"
@@ -177,7 +178,7 @@ class InvokeBaseMixin:
         except Exception as e:
             if ctx.obj and ctx.obj.log_level and ctx.obj.log_level <= logging.DEBUG:
                 # If the user has requested verbose output, print the full traceback
-                console = Console()
+                console = get_console()
                 console.print(Traceback.from_exception(type(e), e, e.__traceback__))
                 exit(1)
             else:
@@ -381,6 +382,11 @@ def format(title: str, vals: Iterable[Any], of: OutputFormat = "table") -> Table
             if not vals:
                 return pretty_repr([])
             return pretty_repr([v.to_dict() for v in vals])
+        case "json-raw":
+            if not vals:
+                return []
+            return json.dumps([v.to_dict() for v in vals])
+
     raise click.ClickException("Unknown output format. Supported formats are: table, table-simple, json.")
 
 
@@ -395,3 +401,10 @@ def get_panel(title: str, renderable: Any, of: OutputFormat = "table") -> Panel:
         title=f"[{PREFERRED_ACCENT_COLOR}]{title}[/{PREFERRED_ACCENT_COLOR}]",
         border_style=PREFERRED_BORDER_COLOR,
     )
+
+
+def get_console() -> Console:
+    """
+    Get a console that is configured to use colors if the terminal supports it.
+    """
+    return Console(color_system="auto", force_terminal=True, width=120)