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

flyte/_task.py

@@ -1,12 +1,12 @@
 from __future__ import annotations
 
+import asyncio
 import weakref
 from dataclasses import dataclass, field, replace
-from functools import cached_property
+from inspect import iscoroutinefunction
 from typing import (
     TYPE_CHECKING,
     Any,
-    Awaitable,
     Callable,
     Coroutine,
     Dict,
@@ -15,36 +15,46 @@ from typing import (
     Literal,
     Optional,
     ParamSpec,
+    Tuple,
+    TypeAlias,
     TypeVar,
     Union,
+    cast,
+    overload,
 )
 
-from flyteidl.core.tasks_pb2 import DataLoadingConfig
-
+from flyte._pod import PodTemplate
 from flyte.errors import RuntimeSystemError, RuntimeUserError
 
 from ._cache import Cache, CacheRequest
 from ._context import internal_ctx
-from ._datastructures import NativeInterface, SerializationContext
 from ._doc import Documentation
 from ._image import Image
+from ._link import Link
 from ._resources import Resources
 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 kubernetes.client import V1PodTemplate
+    from flyteidl2.core.tasks_pb2 import DataLoadingConfig
 
     from ._task_environment import TaskEnvironment
 
 P = ParamSpec("P")  # capture the function's parameters
 R = TypeVar("R")  # return type
 
+AsyncFunctionType: TypeAlias = Callable[P, Coroutine[Any, Any, R]]
+SyncFunctionType: TypeAlias = Callable[P, R]
+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
@@ -64,55 +74,60 @@ 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
     will not be reused across task invocations.
     :param docs: Optional The documentation for the task, if not provided the function docstring will be used.
-    :param env: Optional The environment variables to set for the task.
+    :param env_vars: Optional The environment variables to set for the task.
     :param secrets: Optional The secrets that will be injected into the task at runtime.
     :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
     interface: NativeInterface
+    short_name: str = ""
     task_type: str = "python"
     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, Literal["auto"], None] = None
+    reusable: Union[ReusePolicy, None] = None
     docs: Optional[Documentation] = None
-    env: Optional[Dict[str, str]] = None
+    env_vars: Optional[Dict[str, str]] = None
     secrets: Optional[SecretRequest] = None
     timeout: Optional[TimeoutType] = None
-    primary_container_name: str = "primary"
-    pod_template: Optional[Union[str, V1PodTemplate]] = 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
-    local: bool = field(default=False, init=False)
+    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)
+    links: Tuple[Link, ...] = field(default_factory=tuple)
 
-    def __post_init__(self):
-        # If pod_template is set to a pod, verify
-        if self.pod_template is not None and not isinstance(self.pod_template, str):
-            try:
-                from kubernetes.client import V1PodTemplate  # noqa: F401
-            except ImportError as e:
-                raise ImportError(
-                    "kubernetes is not installed, please install kubernetes package to use pod_template"
-                ) from e
+    # Only used in python 3.10 and 3.11, where we cannot use markcoroutinefunction
+    _call_as_synchronous: bool = False
 
+    def __post_init__(self):
         # Auto set the image based on the image request
         if self.image == "auto":
-            self.image = Image.auto()
+            self.image = Image.from_debian_base()
         elif isinstance(self.image, str):
-            self.image = Image.from_prebuilt(str(self.image))
+            self.image = Image.from_base(str(self.image))
 
         # Auto set cache based on the cache request
         if isinstance(self.cache, str):
@@ -128,6 +143,10 @@ class TaskTemplate(Generic[P, R]):
         if isinstance(self.retries, int):
             self.retries = RetryStrategy(count=self.retries)
 
+        if self.short_name == "":
+            # If short_name is not set, use the name of the task
+            self.short_name = self.name
+
     def __getstate__(self):
         """
         This method is called when the object is pickled. We need to remove the parent_env reference
@@ -145,6 +164,10 @@ class TaskTemplate(Generic[P, R]):
         self.__dict__.update(state)
         self.parent_env = None
 
+    @property
+    def source_file(self) -> Optional[str]:
+        return None
+
     async def pre(self, *args, **kwargs) -> Dict[str, Any]:
         """
         This is the preexecute function that will be
@@ -208,13 +231,78 @@ class TaskTemplate(Generic[P, R]):
     def native_interface(self) -> NativeInterface:
         return self.interface
 
-    async def __call__(self, *args: P.args, **kwargs: P.kwargs) -> Coroutine[Any, Any, R]:
+    @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
+        tasks to be used within an asyncio parent task.
+        This function will also re-raise exceptions from the underlying task.
+
+        Example:
+        ```python
+        @env.task
+        def my_legacy_task(x: int) -> int:
+            return x
+
+        @env.task
+        async def my_new_parent_task(n: int) -> List[int]:
+            collect = []
+            for x in range(n):
+                collect.append(my_legacy_task.aio(x))
+            return asyncio.gather(*collect)
+        ```
+        :param args:
+        :param kwargs:
+        :return:
+        """
+        ctx = internal_ctx()
+        if ctx.is_task_context():
+            from ._internal.controllers import get_controller
+
+            # If we are in a task context, that implies we are executing a Run.
+            # In this scenario, we should submit the task to the controller.
+            controller = get_controller()
+            if controller:
+                if self._call_as_synchronous:
+                    fut = controller.submit_sync(self, *args, **kwargs)
+                    asyncio_future = asyncio.wrap_future(fut)  # Wrap the future to make it awaitable
+                    return await asyncio_future
+                else:
+                    return await controller.submit(self, *args, **kwargs)
+            else:
+                raise RuntimeSystemError("BadContext", "Controller is not initialized.")
+        else:
+            from flyte._logging import logger
+
+            logger.warning(f"Task {self.name} running aio outside of a task context.")
+            # Local execute, just stay out of the way, but because .aio is used, we want to return an awaitable,
+            # 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.
         Please do not override this method, if you simply want to modify the execution behavior, override the
         execute method.
 
-        # TODO lets provide one hook to implement, _pre, _execute and _post. We do not want actual execute to
+        This needs to be overridable to maybe be async.
+        The returned thing from here needs to be an awaitable if the underlying task is async, and a regular object
+        if the task is not.
         """
         try:
             ctx = internal_ctx()
@@ -224,10 +312,19 @@ class TaskTemplate(Generic[P, R]):
                 # 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()
-                if controller:
-                    return await controller.submit(self, *args, **kwargs)
-            return await self.execute(*args, **kwargs)
+                controller = get_controller()
+                if not controller:
+                    raise RuntimeSystemError("BadContext", "Controller is not initialized.")
+
+                if self._call_as_synchronous:
+                    fut = controller.submit_sync(self, *args, **kwargs)
+                    x = fut.result(None)
+                    return x
+                else:
+                    return controller.submit(self, *args, **kwargs)
+            else:
+                # If not in task context, purely function run, stay out of the way
+                return self.forward(*args, **kwargs)
         except RuntimeSystemError:
             raise
         except RuntimeUserError:
@@ -235,33 +332,93 @@ class TaskTemplate(Generic[P, R]):
         except Exception as e:
             raise RuntimeUserError(type(e).__name__, str(e)) from e
 
+    def forward(self, *args: P.args, **kwargs: P.kwargs) -> Coroutine[Any, Any, R] | R:
+        """
+        Think of this as a local execute method for your task. This function will be invoked by the __call__ method
+        when not in a Flyte task execution context.  See the implementation below for an example.
+
+        :param args:
+        :param kwargs:
+        :return:
+        """
+        raise NotImplementedError
+
     def override(
         self,
         *,
-        local: Optional[bool] = None,
-        ref: Optional[bool] = None,
+        short_name: Optional[str] = None,
         resources: Optional[Resources] = None,
-        cache: CacheRequest = "auto",
+        cache: Optional[CacheRequest] = None,
         retries: Union[int, RetryStrategy] = 0,
         timeout: Optional[TimeoutType] = None,
-        reusable: Union[ReusePolicy, Literal["auto"], None] = None,
-        env: Optional[Dict[str, str]] = None,
+        reusable: Union[ReusePolicy, Literal["off"], None] = None,
+        env_vars: Optional[Dict[str, str]] = None,
         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,
+        links: Tuple[Link, ...] = (),
         **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 interruptible: Optional override for the interruptible policy for the task.
+        :param links: Optional override for the Links associated with 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.
         """
-        resources = resources or self.resources
         cache = cache or self.cache
         retries = retries or self.retries
         timeout = timeout or self.timeout
+        max_inline_io_bytes = max_inline_io_bytes or self.max_inline_io_bytes
+
         reusable = reusable or self.reusable
-        env = env or self.env
+        if reusable == "off":
+            reusable = None
+
+        if reusable is not None:
+            if resources is not None:
+                raise ValueError(
+                    "Cannot override resources when reusable is set."
+                    " Reusable tasks will use the parent env's resources. You can disable reusability and"
+                    " override resources if needed. (set reusable='off')"
+                )
+            if env_vars is not None:
+                raise ValueError(
+                    "Cannot override env when reusable is set."
+                    " Reusable tasks will use the parent env's env. You can disable reusability and "
+                    "override env if needed. (set reusable='off')"
+                )
+            if secrets is not None:
+                raise ValueError(
+                    "Cannot override secrets when reusable is set."
+                    " Reusable tasks will use the parent env's secrets. You can disable reusability and "
+                    "override secrets if needed. (set reusable='off')"
+                )
+
+        resources = resources or self.resources
+        env_vars = env_vars or self.env_vars
         secrets = secrets or self.secrets
-        local = local or self.local
-        ref = ref or self.ref
+
+        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")
@@ -271,45 +428,72 @@ class TaskTemplate(Generic[P, R]):
                 raise ValueError("Docs cannot be overridden")
             if k == "interface":
                 raise ValueError("Interface cannot be overridden")
+
         return replace(
             self,
+            short_name=short_name or self.short_name,
             resources=resources,
             cache=cache,
             retries=retries,
             timeout=timeout,
-            reusable=reusable,
-            env=env,
+            reusable=cast(Optional[ReusePolicy], reusable),
+            env_vars=env_vars,
             secrets=secrets,
+            max_inline_io_bytes=max_inline_io_bytes,
+            pod_template=pod_template,
+            interruptible=interruptible,
+            queue=queue or self.queue,
+            links=links or self.links,
+            **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: Callable[P, Awaitable[R]]
+    func: F
+    plugin_config: Optional[Any] = None  # This is used to pass plugin specific configuration
+    debuggable: bool = True
 
-    @cached_property
-    def native_interface(self) -> NativeInterface:
-        return NativeInterface.from_callable(self.func)
+    def __post_init__(self):
+        super().__post_init__()
+        if not iscoroutinefunction(self.func):
+            self._call_as_synchronous = True
+
+    @property
+    def source_file(self) -> Optional[str]:
+        """
+        Returns the source file of the function, if available. This is useful for debugging and tracing.
+        """
+        if hasattr(self.func, "__code__") and self.func.__code__:
+            return self.func.__code__.co_filename
+        return None
+
+    def forward(self, *args: P.args, **kwargs: P.kwargs) -> Coroutine[Any, Any, R] | R:
+        # In local execution, we want to just call the function. Note we're not awaiting anything here.
+        # If the function was a coroutine function, the coroutine is returned and the await that the caller has
+        # in front of the task invocation will handle the awaiting.
+        return self.func(*args, **kwargs)
 
     async def execute(self, *args: P.args, **kwargs: P.kwargs) -> R:
         """
         This is the execute method that will be called when the task is invoked. It will call the actual function.
         # TODO We may need to keep this as the bare func execute, and need a pre and post execute some other func.
         """
+
         ctx = internal_ctx()
+        assert ctx.data.task_context is not None, "Function should have already returned if not in a task context"
         ctx_data = await self.pre(*args, **kwargs)
-        if ctx.data.task_context is not None:
-            tctx = ctx.data.task_context.replace(data=ctx_data)
-            with ctx.replace_task_context(tctx):
+        tctx = ctx.data.task_context.replace(data=ctx_data)
+        with ctx.replace_task_context(tctx):
+            if iscoroutinefunction(self.func):
                 v = await self.func(*args, **kwargs)
-                await self.post(v)
-        else:
-            v = await self.func(*args, **kwargs)
+            else:
+                v = self.func(*args, **kwargs)
             await self.post(v)
         return v
 
@@ -352,6 +536,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,