flyte 0.2.0b7__py3-none-any.whl → 0.2.0b9__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,6 +16,7 @@ from typing import (
     Literal,
     Optional,
     ParamSpec,
+    TypeAlias,
     TypeVar,
     Union,
 )
@@ -42,6 +44,10 @@ if TYPE_CHECKING:
 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]):
@@ -98,6 +104,9 @@ class TaskTemplate(Generic[P, R]):
     local: bool = field(default=False, init=False)
     ref: bool = field(default=False, init=False, repr=False, compare=False)
 
+    # Only used in python 3.10 and 3.11, where we cannot use markcoroutinefunction
+    _call_as_synchronous: bool = 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):
@@ -208,13 +217,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()
@@ -225,9 +281,18 @@ class TaskTemplate(Generic[P, R]):
                 from ._internal.controllers import get_controller
 
                 controller = get_controller()
-                if controller:
-                    return await controller.submit(self, *args, **kwargs)
-            return await self.execute(*args, **kwargs)
+                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 +300,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 +366,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,11 +1,19 @@
 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
 
@@ -23,8 +31,7 @@ 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
@@ -54,7 +61,7 @@ class TaskEnvironment(Environment):
     """
 
     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,
+        **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 env_dep_hints is not None:
+            kwargs["env_dep_hints"] = env_dep_hints
+        return replace(self, **kwargs)
+
+    def task(
         self,
         _func=None,
         *,
@@ -104,6 +121,7 @@ class TaskEnvironment(Environment):
         report: bool = False,
     ) -> Union[AsyncFunctionTaskTemplate, Callable[P, R]]:
         """
+        :param _func: Optional The function to decorate. If not provided, the decorator will return a callable that
         :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.
@@ -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]:
+        def decorator(func: FunctionTypes) -> 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,
+
+            tmpl: AsyncFunctionTaskTemplate = AsyncFunctionTaskTemplate(
+                func=func,
                 name=task_name,
                 image=self.image,
                 resources=self.resources,
@@ -160,29 +165,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/_utils/coro_management.py

@@ -21,5 +21,3 @@ async def run_coros(*coros: typing.Coroutine, return_when: str = asyncio.FIRST_C
         err = t.exception()
         if err:
             raise err
-        else:
-            print(f"Task result: {t.result()}")

flyte/_utils/helpers.py

@@ -1,6 +1,7 @@
 import os
 import string
 import typing
+from contextlib import contextmanager
 from pathlib import Path
 
 
@@ -106,3 +107,17 @@ def get_cwd_editable_install() -> typing.Optional[Path]:
                     return install  # note we want the install folder, not the parent
 
     return None
+
+
+@contextmanager
+def _selector_policy():
+    import asyncio
+
+    original_policy = asyncio.get_event_loop_policy()
+    try:
+        if os.name == "nt" and hasattr(asyncio, "WindowsSelectorEventLoopPolicy"):
+            asyncio.set_event_loop_policy(asyncio.WindowsSelectorEventLoopPolicy())
+
+        yield
+    finally:
+        asyncio.set_event_loop_policy(original_policy)

flyte/_version.py

@@ -17,5 +17,5 @@ __version__: str
 __version_tuple__: VERSION_TUPLE
 version_tuple: VERSION_TUPLE
 
-__version__ = version = '0.2.0b7'
-__version_tuple__ = version_tuple = (0, 2, 0, 'b7')
+__version__ = version = '0.2.0b9'
+__version_tuple__ = version_tuple = (0, 2, 0, 'b9')

flyte/cli/__init__.py

@@ -1,10 +1,3 @@
-"""
-# CLI for Flyte
-
-The flyte cli follows a simple verb based structure, where the top-level commands are verbs that describe the action
-to be taken, and the subcommands are nouns that describe the object of the action.
-"""
-
 from flyte.cli.main import main
 
 __all__ = ["main"]

flyte/cli/_abort.py

@@ -6,7 +6,7 @@ from flyte.cli import _common as common
 @click.group(name="abort")
 def abort():
     """
-    Abort a run.
+    Abort an ongoing process.
     """
 
 

flyte/cli/_common.py

@@ -58,7 +58,7 @@ DRY_RUN_OPTION = click.Option(
 
 def _common_options() -> List[click.Option]:
     """
-    Common options for that will be added to all commands and groups that inherit from CommandBase or GroupBase.
+    Common options that will be added to all commands and groups that inherit from CommandBase or GroupBase.
     """
     return [PROJECT_OPTION, DOMAIN_OPTION]
 
@@ -74,6 +74,7 @@ class CLIConfig:
     """
 
     config: Config
+    ctx: click.Context
     log_level: int | None = logging.ERROR
     endpoint: str | None = None
     insecure: bool = False
@@ -109,8 +110,8 @@ class CLIConfig:
 
 class InvokeBaseMixin:
     """
-    Mixin to catch grpc.RpcError, flyte.RpcError, other errors and other exceptions and raise them as
-     gclick.ClickException.
+    Mixin to catch grpc.RpcError, flyte.RpcError, other errors and other exceptions
+    and raise them as gclick.ClickException.
     """
 
     def invoke(self, ctx):

flyte/cli/_create.py

@@ -1,5 +1,5 @@
 from pathlib import Path
-from typing import get_args
+from typing import Any, Dict, get_args
 
 import rich_click as click
 
@@ -10,7 +10,7 @@ from flyte.remote._secret import SecretTypes
 @click.group(name="create")
 def create():
     """
-    Create a new task or environment.
+    Create resources in a Flyte deployment.
     """
 
 
@@ -32,7 +32,27 @@ def secret(
     domain: str | None = None,
 ):
     """
-    Create a new secret.
+    Create a new secret. The name of the secret is required. For example:
+
+    ```bash
+    $ flyte create secret my_secret --value my_value
+    ```
+
+    If `--from-file` is specified, the value will be read from the file instead of being provided directly:
+
+    ```bash
+    $ flyte create secret my_secret --from-file /path/to/secret_file
+    ```
+
+    The `--type` option can be used to create specific types of secrets.
+    Either `regular` or `image_pull` can be specified.
+    Secrets intended to access container images should be specified as `image_pull`.
+    Other secrets should be specified as `regular`.
+    If no type is specified, `regular` is assumed.
+
+    ```bash
+    $ flyte create secret my_secret --type image_pull
+    ```
     """
     from flyte.remote import Secret
 
@@ -45,19 +65,29 @@ def secret(
 
 @create.command(cls=common.CommandBase)
 @click.option("--endpoint", type=str, help="Endpoint of the Flyte backend.")
-@click.option("--insecure", is_flag=True, help="Use insecure connection to the Flyte backend.")
+@click.option("--insecure", is_flag=True, help="Use an insecure connection to the Flyte backend.")
 @click.option(
     "--org",
     type=str,
     required=False,
-    help="Organization to use, this will override the organization in the config file.",
+    help="Organization to use, this will override the organization in the configusraion file.",
 )
 @click.option(
     "-o",
     "--output",
-    type=click.Path(),
-    default=Path.cwd(),
-    help="Path to the output dir where the config will be saved, defaults to current directory.",
+    type=click.Path(exists=False, writable=True),
+    default=Path.cwd() / "config.yaml",
+    help="Path to the output directory where the configuration will be saved. Defaults to current directory.",
+    show_default=True,
+)
+@click.option(
+    "--force",
+    is_flag=True,
+    default=False,
+    help="Force overwrite of the configuration file if it already exists.",
+    show_default=True,
+    prompt="Are you sure you want to overwrite the configuration file?",
+    confirmation_prompt=True,
 )
 def config(
     output: Path,
@@ -66,25 +96,41 @@ def config(
     org: str | None = None,
     project: str | None = None,
     domain: str | None = None,
+    force: bool = False,
 ):
     """
-    Create a new config file.
+    Creates a configuration file for Flyte CLI.
+    If the `--output` option is not specified, it will create a file named `config.yaml` in the current directory.
+    If the file already exists, it will raise an error unless the `--force` option is used.
     """
     import yaml
 
-    output_file = output / "config.yaml"
-    with open(output_file, "w") as f:
-        d = {
-            "admin": {
-                "endpoint": endpoint,
-                "insecure": insecure,
-            },
-            "task": {
-                "org": org,
-                "project": project,
-                "domain": domain,
-            },
-        }
+    if output.exists() and not force:
+        raise click.BadParameter(f"Output file {output} already exists. Use --force to overwrite.")
+
+    admin: Dict[str, Any] = {}
+    if endpoint:
+        admin["endpoint"] = endpoint
+    if insecure:
+        admin["insecure"] = insecure
+
+    task: Dict[str, str] = {}
+    if org:
+        task["org"] = org
+    if project:
+        task["project"] = project
+    if domain:
+        task["domain"] = domain
+
+    if not admin and not task:
+        raise click.BadParameter("At least one of --endpoint or --org must be provided.")
+
+    with open(output, "w") as f:
+        d: Dict[str, Any] = {}
+        if admin:
+            d["admin"] = admin
+        if task:
+            d["task"] = task
         yaml.dump(d, f)
 
-    click.echo(f"Config file created at {output_file}")
+    click.echo(f"Config file created at {output}")

flyte/cli/_delete.py

@@ -6,7 +6,7 @@ import flyte.cli._common as common
 @click.group(name="delete")
 def delete():
     """
-    Delete a task or environment.
+    Remove resources from a Flyte deployment.
     """
 
 
@@ -15,7 +15,7 @@ def delete():
 @click.pass_obj
 def secret(cfg: common.CLIConfig, name: str, project: str | None = None, domain: str | None = None):
     """
-    Delete a secret.
+    Delete a secret. The name of the secret is required.
     """
     from flyte.remote import Secret
 

flyte/cli/_deploy.py

@@ -61,7 +61,7 @@ class DeployArguments:
     @classmethod
     def options(cls) -> List[click.Option]:
         """
-        Return the set of base parameters added to every pyflyte run workflow subcommand.
+        Return the set of base parameters added to every flyte run workflow subcommand.
         """
         return [common.get_option_from_metadata(f.metadata) for f in fields(cls) if f.metadata]
 
@@ -87,7 +87,7 @@ class DeployEnvCommand(click.Command):
 
 class EnvPerFileGroup(common.ObjectsPerFileGroup):
     """
-    Group that creates a command for each task in the current directory that is not __init__.py.
+    Group that creates a command for each task in the current directory that is not `__init__.py`.
     """
 
     def __init__(self, filename: Path, deploy_args: DeployArguments, *args, **kwargs):
@@ -111,7 +111,7 @@ class EnvPerFileGroup(common.ObjectsPerFileGroup):
 
 class EnvFiles(common.FileGroup):
     """
-    Group that creates a command for each file in the current directory that is not __init__.py.
+    Group that creates a command for each file in the current directory that is not `__init__.py`.
     """
 
     common_options_enabled = False
@@ -138,5 +138,8 @@ class EnvFiles(common.FileGroup):
 
 deploy = EnvFiles(
     name="deploy",
-    help="deploy one or more environments from a python file.",
+    help="""
+    Deploy one or more environments from a python file.
+    The deploy command will create or update environments in the Flyte system.
+    """,
 )