flyte 0.0.1b3__py3-none-any.whl → 0.2.0a0__py3-none-any.whl

flyte/_task.py

@@ -1,12 +1,13 @@
 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,17 +16,18 @@ from typing import (
     Literal,
     Optional,
     ParamSpec,
+    TypeAlias,
     TypeVar,
     Union,
 )
 
 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 ._resources import Resources
@@ -33,15 +35,18 @@ from ._retry import RetryStrategy
 from ._reusable_environment import ReusePolicy
 from ._secret import SecretRequest
 from ._timeout import TimeoutType
+from .models import NativeInterface, SerializationContext
 
 if TYPE_CHECKING:
-    from kubernetes.client import V1PodTemplate
-
     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 = Union[AsyncFunctionType, SyncFunctionType]
+
 
 @dataclass(kw_only=True)
 class TaskTemplate(Generic[P, R]):
@@ -78,6 +83,7 @@ class TaskTemplate(Generic[P, R]):
 
     name: str
     interface: NativeInterface
+    friendly_name: str = ""
     task_type: str = "python"
     task_type_version: int = 0
     image: Union[str, Image, Literal["auto"]] = "auto"
@@ -90,24 +96,17 @@ class TaskTemplate(Generic[P, R]):
     env: 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
 
     parent_env: Optional[weakref.ReferenceType[TaskEnvironment]] = None
     local: bool = field(default=False, init=False)
     ref: bool = field(default=False, init=False, repr=False, compare=False)
 
-    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()
@@ -128,6 +127,10 @@ class TaskTemplate(Generic[P, R]):
         if isinstance(self.retries, int):
             self.retries = RetryStrategy(count=self.retries)
 
+        if self.friendly_name == "":
+            # If friendly_name is not set, use the name of the task
+            self.friendly_name = self.name
+
     def __getstate__(self):
         """
         This method is called when the object is pickled. We need to remove the parent_env reference
@@ -208,13 +211,60 @@ 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]:
+    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:
+            # 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)
+
+    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 +274,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,6 +294,17 @@ 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,
         *,
@@ -290,26 +360,38 @@ class AsyncFunctionTaskTemplate(TaskTemplate[P, R]):
     is decorated with the task decorator.
     """
 
-    func: Callable[P, Awaitable[R]]
+    func: FunctionTypes
+
+    def __post_init__(self):
+        super().__post_init__()
+        if not iscoroutinefunction(self.func):
+            self._call_as_synchronous = True
 
     @cached_property
     def native_interface(self) -> NativeInterface:
         return NativeInterface.from_callable(self.func)
 
+    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
 

flyte/_task_environment.py

@@ -1,16 +1,23 @@
 from __future__ import annotations
 
-import asyncio
 import weakref
 from dataclasses import dataclass, field, replace
 from datetime import timedelta
-from functools import wraps
-from typing import TYPE_CHECKING, Awaitable, Callable, Dict, List, Literal, Optional, ParamSpec, TypeVar, Union, cast
+from typing import (
+    TYPE_CHECKING,
+    Any,
+    Callable,
+    Dict,
+    List,
+    Literal,
+    Optional,
+    Union,
+    cast,
+)
 
 import rich.repr
 
 from ._cache import CacheRequest
-from ._datastructures import NativeInterface
 from ._doc import Documentation
 from ._environment import Environment
 from ._image import Image
@@ -19,12 +26,12 @@ from ._retry import RetryStrategy
 from ._reusable_environment import ReusePolicy
 from ._secret import SecretRequest
 from ._task import AsyncFunctionTaskTemplate, TaskTemplate
+from .models import NativeInterface
 
 if TYPE_CHECKING:
     from kubernetes.client import V1PodTemplate
 
-P = ParamSpec("P")  # capture the function's parameters
-R = TypeVar("R")  # return type
+    from ._task import FunctionTypes, P, R
 
 
 @rich.repr.auto
@@ -47,14 +54,14 @@ class TaskEnvironment(Environment):
     :param resources: Resources to allocate for the environment.
     :param env: Environment variables to set for the environment.
     :param secrets: Secrets to inject into the environment.
-    :param env_dep_hints: Environment dependencies to hint, so when you deploy the environment, the dependencies are
+    :param depends_on: 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 cache: Cache policy for the environment.
     :param reusable: Reuse policy for the environment, if set, a python process may be reused for multiple tasks.
     """
 
     cache: Union[CacheRequest] = "auto"
-    reusable: Union[ReusePolicy, Literal["auto"], None] = None
+    reusable: ReusePolicy | None = 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
 
@@ -65,32 +72,42 @@ class TaskEnvironment(Environment):
         name: str,
         image: Optional[Union[str, Image, Literal["auto"]]] = None,
         resources: Optional[Resources] = None,
-        cache: Union[CacheRequest, None] = None,
         env: Optional[Dict[str, str]] = None,
-        reusable: Union[ReusePolicy, None] = None,
         secrets: Optional[SecretRequest] = None,
-        env_dep_hints: Optional[List[Environment]] = None,
+        depends_on: Optional[List[Environment]] = None,
+        **kwargs: Any,
     ) -> TaskEnvironment:
         """
-        Clone the environment with new settings.
+        Clone the TaskEnvironment with new parameters.
+        besides the base environment parameters, you can override, kwargs like `cache`, `reusable`, etc.
+
         """
-        if image is None:
-            image = self.image
-        else:
-            image = "auto"
-        return replace(
-            self,
-            cache=cache if cache else "auto",
-            reusable=reusable,
-            name=name,
-            image=image,
-            resources=resources,
-            env=env,
-            secrets=secrets,
-            env_dep_hints=env_dep_hints if env_dep_hints else [],
-        )
-
-    def _task(
+        cache = kwargs.pop("cache", None)
+        reusable = kwargs.pop("reusable", None)
+
+        # validate unknown kwargs if needed
+        if kwargs:
+            raise TypeError(f"Unexpected keyword arguments: {list(kwargs.keys())}")
+
+        kwargs = self._get_kwargs()
+        kwargs["name"] = name
+        if image is not None:
+            kwargs["image"] = image
+        if resources is not None:
+            kwargs["resources"] = resources
+        if cache is not None:
+            kwargs["cache"] = cache
+        if env is not None:
+            kwargs["env"] = env
+        if reusable is not None:
+            kwargs["reusable"] = reusable
+        if secrets is not None:
+            kwargs["secrets"] = secrets
+        if depends_on is not None:
+            kwargs["depends_on"] = depends_on
+        return replace(self, **kwargs)
+
+    def task(
         self,
         _func=None,
         *,
@@ -104,7 +121,8 @@ class TaskEnvironment(Environment):
         report: bool = False,
     ) -> Union[AsyncFunctionTaskTemplate, Callable[P, R]]:
         """
-        :param name: Optional The name of the task (defaults to the function name)
+        :param _func: Optional The function to decorate. If not provided, the decorator will return a callable that
+        :param 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.
@@ -119,25 +137,12 @@ class TaskEnvironment(Environment):
             if pod_template is not None:
                 raise ValueError("Cannot set pod_template when environment is reusable.")
 
-        def decorator(func: Callable[P, Awaitable[R]]) -> AsyncFunctionTaskTemplate[P, R]:
-            task_name = name or func.__name__
-            task_name = self.name + "." + task_name
-            if len(task_name) > 30:
-                # delete this if we can remove the restriction that task names must be <= 30 characters
-                task_name = task_name[-30:]
-
-            @wraps(func)
-            async def wrapper(*args: P.args, **kwargs: P.kwargs) -> R:
-                return await func(*args, **kwargs)
-
-            if not asyncio.iscoroutinefunction(func):
-                raise TypeError(
-                    f"Function {func.__name__} is not a coroutine function. Use @env.task decorator for async tasks."
-                    f"You can simply mark your function as async def {func.__name__} to make it a coroutine function, "
-                    f"it is ok to write sync code in async functions, but not the other way around."
-                )
-            tmpl = AsyncFunctionTaskTemplate(
-                func=wrapper,
+        def decorator(func: FunctionTypes) -> AsyncFunctionTaskTemplate[P, R]:
+            friendly_name = name or func.__name__
+            task_name = self.name + "." + func.__name__
+
+            tmpl: AsyncFunctionTaskTemplate = AsyncFunctionTaskTemplate(
+                func=func,
                 name=task_name,
                 image=self.image,
                 resources=self.resources,
@@ -152,6 +157,7 @@ class TaskEnvironment(Environment):
                 parent_env=weakref.ref(self),
                 interface=NativeInterface.from_callable(func),
                 report=report,
+                friendly_name=friendly_name,
             )
             self._tasks[task_name] = tmpl
             return tmpl
@@ -160,29 +166,6 @@ class TaskEnvironment(Environment):
             return cast(AsyncFunctionTaskTemplate, decorator)
         return cast(AsyncFunctionTaskTemplate, decorator(_func))
 
-    @property
-    def task(self) -> Callable:
-        """
-        Decorator to create a new task with the environment settings.
-        The task will be executed in its own container with the specified image, resources, and environment variables,
-        unless reusePolicy is set, in which case the same container will be reused for all tasks with the same
-        environment settings.
-
-        :param name: Optional The name of 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.
-        :param docs: Optional The documentation for the task, if not provided the function docstring will be used.
-        :param secrets: Optional The secrets that will be injected into the task at runtime.
-        :param timeout: Optional The timeout for the task.
-        :param pod_template: Optional The pod template for the task, if not provided the default pod template will be
-         used.
-        :param report: Optional Whether to generate the html report for the task, defaults to False.
-
-        :return: New Task instance or Task decorator
-        """
-        return self._task
-
     @property
     def tasks(self) -> Dict[str, TaskTemplate]:
         """

flyte/_trace.py

@@ -4,7 +4,7 @@ import time
 from datetime import timedelta
 from typing import Any, AsyncGenerator, AsyncIterator, Awaitable, Callable, TypeGuard, TypeVar, Union, cast
 
-from flyte._datastructures import NativeInterface
+from flyte.models import NativeInterface
 
 T = TypeVar("T")
 
@@ -14,7 +14,6 @@ def trace(func: Callable[..., T]) -> Callable[..., T]:
     A decorator that traces function execution with timing information.
     Works with regular functions, async functions, and async generators/iterators.
     """
-    func_name = func.__name__
 
     @functools.wraps(func)
     def wrapper_sync(*args: Any, **kwargs: Any) -> Any:
@@ -31,9 +30,9 @@ def trace(func: Callable[..., T]) -> Callable[..., T]:
             # 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()
+            controller = get_controller()
             iface = NativeInterface.from_callable(func)
-            info, ok = await controller.get_action_outputs(iface, func_name, *args, **kwargs)
+            info, ok = await controller.get_action_outputs(iface, func, *args, **kwargs)
             if ok:
                 if info.output:
                     return info.output
@@ -74,9 +73,9 @@ def trace(func: Callable[..., T]) -> Callable[..., T]:
             # 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()
+            controller = get_controller()
             iface = NativeInterface.from_callable(func)
-            info, ok = await controller.get_action_outputs(iface, func_name, *args, **kwargs)
+            info, ok = await controller.get_action_outputs(iface, func, *args, **kwargs)
             if ok:
                 if info.output:
                     for item in info.output:
@@ -88,17 +87,12 @@ def trace(func: Callable[..., T]) -> Callable[..., T]:
                 items = []
                 result = func(*args, **kwargs)
                 # TODO ideally we should use streaming into the type-engine so that it stream uploads large blocks
-                if inspect.isasyncgen(result):
+                if inspect.isasyncgen(result) or is_async_iterable(result):
                     # If it's directly an async generator
                     async_iter = result
                     async for item in async_iter:
                         items.append(item)
                         yield item
-                elif is_async_iterable(result):
-                    # If it's an async iterable (has __aiter__)
-                    async for item in result:
-                        items.append(item)
-                        yield item
                 duration = time.time() - start_time
                 info.add_outputs(items, timedelta(seconds=duration))
                 await controller.record_trace(info)
@@ -118,10 +112,8 @@ def trace(func: Callable[..., T]) -> Callable[..., T]:
     # Choose the appropriate wrapper based on the function type
     if inspect.iscoroutinefunction(func):
         # This handles async functions that return normal values
-        print(f"Coroutine function {func.__name__}")
         return cast(Callable[..., T], wrapper_async)
     elif inspect.isasyncgenfunction(func):
-        print(f"Async generator function {func.__name__}")
         return cast(Callable[..., T], wrapper_async_iterator)
     else:
         # For regular sync functions

flyte/_utils/__init__.py

@@ -4,17 +4,23 @@ Internal utility functions.
 Except for logging, modules in this package should not depend on any other part of the repo.
 """
 
+from .async_cache import AsyncLRUCache
 from .coro_management import run_coros
 from .file_handling import filehash_update, update_hasher_for_source
 from .helpers import get_cwd_editable_install
 from .lazy_module import lazy_module
+from .org_discovery import hostname_from_url, org_from_endpoint, sanitize_endpoint
 from .uv_script_parser import parse_uv_script_file
 
 __all__ = [
+    "AsyncLRUCache",
     "filehash_update",
     "get_cwd_editable_install",
+    "hostname_from_url",
     "lazy_module",
+    "org_from_endpoint",
     "parse_uv_script_file",
     "run_coros",
+    "sanitize_endpoint",
     "update_hasher_for_source",
 ]