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

flyte/_run.py

@@ -90,9 +90,10 @@ class _Runner:
         env_vars: Dict[str, str] | None = None,
         labels: Dict[str, str] | None = None,
         annotations: Dict[str, str] | None = None,
-        interruptible: bool = False,
+        interruptible: bool | None = None,
         log_level: int | None = None,
         disable_run_cache: bool = False,
+        queue: Optional[str] = None,
     ):
         from flyte._tools import ipython_check
 
@@ -111,8 +112,8 @@ class _Runner:
         self._copy_bundle_to = copy_bundle_to
         self._interactive_mode = interactive_mode if interactive_mode else ipython_check()
         self._raw_data_path = raw_data_path
-        self._metadata_path = metadata_path or "/tmp"
-        self._run_base_dir = run_base_dir or "/tmp/base"
+        self._metadata_path = metadata_path
+        self._run_base_dir = run_base_dir
         self._overwrite_cache = overwrite_cache
         self._project = project
         self._domain = domain
@@ -122,6 +123,7 @@ class _Runner:
         self._interruptible = interruptible
         self._log_level = log_level
         self._disable_run_cache = disable_run_cache
+        self._queue = queue
 
     @requires_initialization
     async def _run_remote(self, obj: TaskTemplate[P, R] | LazyEntity, *args: P.args, **kwargs: P.kwargs) -> Run:
@@ -150,6 +152,7 @@ class _Runner:
             version = task.pb2.task_id.version
             code_bundle = None
         else:
+            task = cast(TaskTemplate[P, R], obj)
             if obj.parent_env is None:
                 raise ValueError("Task is not attached to an environment. Please attach the task to an environment")
 
@@ -204,10 +207,11 @@ class _Runner:
             inputs = await convert_from_native_to_inputs(obj.native_interface, *args, **kwargs)
 
         env = self._env_vars or {}
-        if self._log_level:
-            env["LOG_LEVEL"] = str(self._log_level)
-        else:
-            env["LOG_LEVEL"] = str(logger.getEffectiveLevel())
+        if env.get("LOG_LEVEL") is None:
+            if self._log_level:
+                env["LOG_LEVEL"] = str(self._log_level)
+            else:
+                env["LOG_LEVEL"] = str(logger.getEffectiveLevel())
 
         if not self._dry_run:
             if get_client() is None:
@@ -263,10 +267,13 @@ class _Runner:
                         inputs=inputs.proto_inputs,
                         run_spec=run_definition_pb2.RunSpec(
                             overwrite_cache=self._overwrite_cache,
-                            interruptible=wrappers_pb2.BoolValue(value=self._interruptible),
+                            interruptible=wrappers_pb2.BoolValue(value=self._interruptible)
+                            if self._interruptible is not None
+                            else None,
                             annotations=annotations,
                             labels=labels,
                             envs=env_kv,
+                            cluster=self._queue or task.queue,
                         ),
                     ),
                 )
@@ -385,6 +392,7 @@ class _Runner:
                 " flyte.with_runcontext(run_base_dir='s3://bucket/metadata/outputs')",
             )
         output_path = self._run_base_dir
+        run_base_dir = self._run_base_dir
         raw_data_path = f"{output_path}/rd/{random_id}"
         raw_data_path_obj = RawDataPath(path=raw_data_path)
         checkpoint_path = f"{raw_data_path}/checkpoint"
@@ -401,7 +409,7 @@ class _Runner:
                 version=version if version else "na",
                 raw_data_path=raw_data_path_obj,
                 compiled_image_cache=image_cache,
-                run_base_dir=self._run_base_dir,
+                run_base_dir=run_base_dir,
                 report=flyte.report.Report(name=action.name),
             )
             async with ctx.replace_task_context(tctx):
@@ -426,6 +434,18 @@ class _Runner:
         else:
             action = ActionID(name=self._name)
 
+        metadata_path = self._metadata_path
+        if metadata_path is None:
+            metadata_path = pathlib.Path("/") / "tmp" / "flyte" / "metadata" / action.name
+        else:
+            metadata_path = pathlib.Path(metadata_path) / action.name
+        output_path = metadata_path / "a0"
+        if self._raw_data_path is None:
+            path = pathlib.Path("/") / "tmp" / "flyte" / "raw_data" / action.name
+            raw_data_path = RawDataPath(path=str(path))
+        else:
+            raw_data_path = RawDataPath(path=self._raw_data_path)
+
         ctx = internal_ctx()
         tctx = TaskContext(
             action=action,
@@ -434,10 +454,10 @@ class _Runner:
                 checkpoint_path=internal_ctx().raw_data.path,
             ),
             code_bundle=None,
-            output_path=self._metadata_path,
-            run_base_dir=self._metadata_path,
+            output_path=str(output_path),
+            run_base_dir=str(metadata_path),
             version="na",
-            raw_data_path=internal_ctx().raw_data,
+            raw_data_path=raw_data_path,
             compiled_image_cache=None,
             report=Report(name=action.name),
             mode="local",
@@ -469,7 +489,7 @@ class _Runner:
 
             @property
             def url(self) -> str:
-                return "local-run"
+                return str(metadata_path)
 
             def wait(
                 self,
@@ -550,9 +570,10 @@ def with_runcontext(
     env_vars: Dict[str, str] | None = None,
     labels: Dict[str, str] | None = None,
     annotations: Dict[str, str] | None = None,
-    interruptible: bool = False,
+    interruptible: bool | None = None,
     log_level: int | None = None,
     disable_run_cache: bool = False,
+    queue: Optional[str] = None,
 ) -> _Runner:
     """
     Launch a new run with the given parameters as the context.
@@ -590,15 +611,20 @@ def with_runcontext(
     :param env_vars: Optional Environment variables to set for the run
     :param labels: Optional Labels to set for the run
     :param annotations: Optional Annotations to set for the run
-    :param interruptible: Optional If true, the run can be interrupted by the user.
+    :param interruptible: Optional If true, the run can be scheduled on interruptible instances and false implies
+        that all tasks in the run should only be scheduled on non-interruptible instances. If not specified the
+        original setting on all tasks is retained.
     :param log_level: Optional Log level to set for the run. If not provided, it will be set to the default log level
         set using `flyte.init()`
     :param disable_run_cache: Optional If true, the run cache will be disabled. This is useful for testing purposes.
+    :param queue: Optional The queue to use for the run. This is used to specify the cluster to use for the run.
 
     :return: runner
     """
     if mode == "hybrid" and not name and not run_base_dir:
         raise ValueError("Run name and run base dir are required for hybrid mode")
+    if copy_style == "none" and not version:
+        raise ValueError("Version is required when copy_style is 'none'")
     return _Runner(
         force_mode=mode,
         name=name,
@@ -619,6 +645,7 @@ def with_runcontext(
         domain=domain,
         log_level=log_level,
         disable_run_cache=disable_run_cache,
+        queue=queue,
     )
 
 

flyte/_task.py

@@ -15,6 +15,7 @@ from typing import (
     Literal,
     Optional,
     ParamSpec,
+    Tuple,
     TypeAlias,
     TypeVar,
     Union,
@@ -38,6 +39,8 @@ from .models import MAX_INLINE_IO_BYTES, NativeInterface, SerializationContext
 if TYPE_CHECKING:
     from flyteidl.core.tasks_pb2 import DataLoadingConfig
 
+    from flyte.trigger import Trigger
+
     from ._task_environment import TaskEnvironment
 
 P = ParamSpec("P")  # capture the function's parameters
@@ -69,8 +72,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 +84,9 @@ 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.
     """
 
     name: str
@@ -90,8 +96,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 +106,12 @@ class TaskTemplate(Generic[P, R]):
     timeout: Optional[TimeoutType] = None
     pod_template: Optional[Union[str, PodTemplate]] = None
     report: bool = False
+    queue: Optional[str] = None
 
     parent_env: Optional[weakref.ReferenceType[TaskEnvironment]] = 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
@@ -327,11 +335,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 +393,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,6 +417,8 @@ 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,
         )
 

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,21 +144,28 @@ 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(
         self,
-        _func=None,
+        _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 = (),
+    ) -> 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