flyte 0.2.0b8__py3-none-any.whl → 0.2.0b10__py3-none-any.whl

flyte/_run.py

@@ -1,9 +1,11 @@
 from __future__ import annotations
 
+import asyncio
 import pathlib
 import uuid
 from typing import TYPE_CHECKING, Any, Literal, Optional, Tuple, Union, cast
 
+import flyte.errors
 from flyte.errors import InitializationError
 from flyte.models import ActionID, Checkpoints, CodeBundle, RawDataPath, SerializationContext, TaskContext
 from flyte.syncify import syncify
@@ -24,6 +26,7 @@ from ._tools import ipython_check
 
 if TYPE_CHECKING:
     from flyte.remote import Run
+    from flyte.remote._task import LazyEntity
 
     from ._code_bundle import CopyFiles
 
@@ -80,8 +83,11 @@ class _Runner:
         self._run_base_dir = run_base_dir or "/tmp/base"
 
     @requires_initialization
-    async def _run_remote(self, obj: TaskTemplate[P, R], *args: P.args, **kwargs: P.kwargs) -> Run:
+    async def _run_remote(self, obj: TaskTemplate[P, R] | LazyEntity, *args: P.args, **kwargs: P.kwargs) -> Run:
+        import grpc
+
         from flyte.remote import Run
+        from flyte.remote._task import LazyEntity
 
         from ._code_bundle import build_code_bundle, build_pkl_bundle
         from ._deploy import build_images, plan_deploy
@@ -92,37 +98,47 @@ class _Runner:
 
         cfg = get_common_config()
 
-        if obj.parent_env is None:
-            raise ValueError("Task is not attached to an environment. Please attach the task to an environment.")
+        if isinstance(obj, LazyEntity):
+            task = await obj.fetch.aio()
+            task_spec = task.pb2.spec
+            inputs = await convert_from_native_to_inputs(task.interface, *args, **kwargs)
+            version = task.pb2.task_id.version
+            code_bundle = None
+        else:
+            if obj.parent_env is None:
+                raise ValueError("Task is not attached to an environment. Please attach the task to an environment")
 
-        deploy_plan = plan_deploy(cast(Environment, obj.parent_env()))
-        image_cache = await build_images(deploy_plan)
+            deploy_plan = plan_deploy(cast(Environment, obj.parent_env()))
+            image_cache = await build_images(deploy_plan)
 
-        if self._interactive_mode:
-            code_bundle = await build_pkl_bundle(
-                obj, upload_to_controlplane=not self._dry_run, copy_bundle_to=self._copy_bundle_to
-            )
-        else:
-            if self._copy_files != "none":
-                code_bundle = await build_code_bundle(
-                    from_dir=cfg.root_dir, dryrun=self._dry_run, copy_bundle_to=self._copy_bundle_to
+            if self._interactive_mode:
+                code_bundle = await build_pkl_bundle(
+                    obj, upload_to_controlplane=not self._dry_run, copy_bundle_to=self._copy_bundle_to
                 )
             else:
-                code_bundle = None
+                if self._copy_files != "none":
+                    code_bundle = await build_code_bundle(
+                        from_dir=cfg.root_dir,
+                        dryrun=self._dry_run,
+                        copy_bundle_to=self._copy_bundle_to,
+                        copy_style=self._copy_files,
+                    )
+                else:
+                    code_bundle = None
 
-        version = self._version or (
-            code_bundle.computed_version if code_bundle and code_bundle.computed_version else None
-        )
-        if not version:
-            raise ValueError("Version is required when running a task")
-        s_ctx = SerializationContext(
-            code_bundle=code_bundle,
-            version=version,
-            image_cache=image_cache,
-            root_dir=cfg.root_dir,
-        )
-        task_spec = translate_task_to_wire(obj, s_ctx)
-        inputs = await convert_from_native_to_inputs(obj.native_interface, *args, **kwargs)
+            version = self._version or (
+                code_bundle.computed_version if code_bundle and code_bundle.computed_version else None
+            )
+            if not version:
+                raise ValueError("Version is required when running a task")
+            s_ctx = SerializationContext(
+                code_bundle=code_bundle,
+                version=version,
+                image_cache=image_cache,
+                root_dir=cfg.root_dir,
+            )
+            task_spec = translate_task_to_wire(obj, s_ctx)
+            inputs = await convert_from_native_to_inputs(obj.native_interface, *args, **kwargs)
 
         if not self._dry_run:
             if get_client() is None:
@@ -159,15 +175,35 @@ class _Runner:
             if task_spec.task_template.id.version == "":
                 task_spec.task_template.id.version = version
 
-            resp = await get_client().run_service.CreateRun(
-                run_service_pb2.CreateRunRequest(
-                    run_id=run_id,
-                    project_id=project_id,
-                    task_spec=task_spec,
-                    inputs=inputs.proto_inputs,
-                ),
-            )
-            return Run(pb2=resp.run)
+            try:
+                resp = await get_client().run_service.CreateRun(
+                    run_service_pb2.CreateRunRequest(
+                        run_id=run_id,
+                        project_id=project_id,
+                        task_spec=task_spec,
+                        inputs=inputs.proto_inputs,
+                    ),
+                )
+                return Run(pb2=resp.run)
+            except grpc.aio.AioRpcError as e:
+                if e.code() == grpc.StatusCode.UNAVAILABLE:
+                    raise flyte.errors.RuntimeSystemError(
+                        "SystemUnavailableError",
+                        "Flyte system is currently unavailable. check your configuration, or the service status.",
+                    ) from e
+                elif e.code() == grpc.StatusCode.INVALID_ARGUMENT:
+                    raise flyte.errors.RuntimeUserError("InvalidArgumentError", e.details())
+                elif e.code() == grpc.StatusCode.ALREADY_EXISTS:
+                    # TODO maybe this should be a pass and return existing run?
+                    raise flyte.errors.RuntimeUserError(
+                        "RunAlreadyExistsError",
+                        f"A run with the name '{self._name}' already exists. Please choose a different name.",
+                    )
+                else:
+                    raise flyte.errors.RuntimeSystemError(
+                        "RunCreationError",
+                        f"Failed to create run: {e.details()}",
+                    ) from e
 
         class DryRun(Run):
             def __init__(self, _task_spec, _inputs, _code_bundle):
@@ -224,7 +260,10 @@ class _Runner:
             else:
                 if self._copy_files != "none":
                     code_bundle = await build_code_bundle(
-                        from_dir=cfg.root_dir, dryrun=self._dry_run, copy_bundle_to=self._copy_bundle_to
+                        from_dir=cfg.root_dir,
+                        dryrun=self._dry_run,
+                        copy_bundle_to=self._copy_bundle_to,
+                        copy_style=self._copy_files,
                     )
                 else:
                     code_bundle = None
@@ -287,43 +326,44 @@ class _Runner:
 
     async def _run_local(self, obj: TaskTemplate[P, R], *args: P.args, **kwargs: P.kwargs) -> R:
         from flyte._internal.controllers import create_controller
-        from flyte._internal.runtime.convert import (
-            convert_error_to_native,
-            convert_from_native_to_inputs,
-            convert_outputs_to_native,
-        )
-        from flyte._internal.runtime.entrypoints import direct_dispatch
+        from flyte._internal.controllers._local_controller import LocalController
+        from flyte.report import Report
+
+        controller = cast(LocalController, create_controller("local"))
 
-        controller = create_controller("local")
-        inputs = await convert_from_native_to_inputs(obj.native_interface, *args, **kwargs)
         if self._name is None:
             action = ActionID.create_random()
         else:
             action = ActionID(name=self._name)
-        out, err = await direct_dispatch(
-            obj,
+
+        ctx = internal_ctx()
+        tctx = TaskContext(
             action=action,
-            raw_data_path=internal_ctx().raw_data,
-            version="na",
-            controller=controller,
-            inputs=inputs,
-            output_path=self._metadata_path,
-            run_base_dir=self._metadata_path,
             checkpoints=Checkpoints(
                 prev_checkpoint_path=internal_ctx().raw_data.path, checkpoint_path=internal_ctx().raw_data.path
             ),
-        )  # type: ignore
-        if err:
-            native_err = convert_error_to_native(err)
-            if native_err:
-                raise native_err
-        if obj.native_interface.outputs and len(obj.native_interface.outputs) > 0:
-            if out is not None:
-                return cast(R, await convert_outputs_to_native(obj.native_interface, out))
-        return cast(R, None)
+            code_bundle=None,
+            output_path=self._metadata_path,
+            run_base_dir=self._metadata_path,
+            version="na",
+            raw_data_path=internal_ctx().raw_data,
+            compiled_image_cache=None,
+            report=Report(name=action.name),
+            mode="local",
+        )
+        with ctx.replace_task_context(tctx):
+            # make the local version always runs on a different thread, returns a wrapped future.
+            if obj._call_as_synchronous:
+                fut = controller.submit_sync(obj, *args, **kwargs)
+                awaitable = asyncio.wrap_future(fut)
+                return await awaitable
+            else:
+                return await controller.submit(obj, *args, **kwargs)
 
     @syncify
-    async def run(self, task: TaskTemplate[P, Union[R, Run]], *args: P.args, **kwargs: P.kwargs) -> Union[R, Run]:
+    async def run(
+        self, task: TaskTemplate[P, Union[R, Run]] | LazyEntity, *args: P.args, **kwargs: P.kwargs
+    ) -> Union[R, Run]:
         """
         Run an async `@env.task` or `TaskTemplate` instance. The existing async context will be used.
 
@@ -345,13 +385,18 @@ class _Runner:
         :param kwargs: Keyword arguments to pass to the Task
         :return: Run instance or the result of the task
         """
+        from flyte.remote._task import LazyEntity
+
+        if isinstance(task, LazyEntity) and self._mode != "remote":
+            raise ValueError("Remote task can only be run in remote mode.")
         if self._mode == "remote":
             return await self._run_remote(task, *args, **kwargs)
+        task = cast(TaskTemplate, task)
         if self._mode == "hybrid":
             return await self._run_hybrid(task, *args, **kwargs)
 
         # TODO We could use this for remote as well and users could simply pass flyte:// or s3:// or file://
-        async with internal_ctx().new_raw_data_path(
+        with internal_ctx().new_raw_data_path(
             raw_data_path=RawDataPath.from_local_folder(local_folder=self._raw_data_path)
         ):
             return await self._run_local(task, *args, **kwargs)

flyte/_task.py

@@ -1,5 +1,6 @@
 from __future__ import annotations
 
+import asyncio
 import weakref
 from dataclasses import dataclass, field, replace
 from functools import cached_property
@@ -216,6 +217,51 @@ class TaskTemplate(Generic[P, R]):
     def native_interface(self) -> NativeInterface:
         return self.interface
 
+    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.
@@ -235,15 +281,18 @@ class TaskTemplate(Generic[P, R]):
                 from ._internal.controllers import get_controller
 
                 controller = get_controller()
-                if controller:
-                    if self._call_as_synchronous:
-                        return controller.submit_sync(self, *args, **kwargs)
-                    else:
-                        return controller.submit(self, *args, **kwargs)
-                else:
+                if not controller:
                     raise RuntimeSystemError("BadContext", "Controller is not initialized.")
 
-            return self.forward(*args, **kwargs)
+                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:

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.0b8'
-__version_tuple__ = version_tuple = (0, 2, 0, 'b8')
+__version__ = version = '0.2.0b10'
+__version_tuple__ = version_tuple = (0, 2, 0, 'b10')

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

@@ -32,7 +32,7 @@ PROJECT_OPTION = click.Option(
     required=False,
     type=str,
     default=None,
-    help="Project to operate on",
+    help="Project to which this command applies.",
     show_default=True,
 )
 
@@ -41,7 +41,7 @@ DOMAIN_OPTION = click.Option(
     required=False,
     type=str,
     default=None,
-    help="Domain to operate on",
+    help="Domain to which this command applies.",
     show_default=True,
 )
 
@@ -51,14 +51,14 @@ DRY_RUN_OPTION = click.Option(
     type=bool,
     is_flag=True,
     default=False,
-    help="Dry run, do not actually call the backend service.",
+    help="Dry run. Do not actually call the backend service.",
     show_default=True,
 )
 
 
 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]
 
@@ -105,13 +105,13 @@ class CLIConfig:
         updated_config = self.config.with_params(platform_cfg, task_cfg)
 
         logger.debug(f"Initializing CLI with config: {updated_config}")
-        flyte.init_auto_from_config(updated_config)
+        flyte.init_from_config(updated_config)
 
 
 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():
     """
-    The create subcommand allows you to create resources in a Flyte deployment.
+    Create resources in a Flyte deployment.
     """
 
 
@@ -32,22 +32,26 @@ def secret(
     domain: str | None = None,
 ):
     """
-    Create a new secret, the name of the secret is required.
+    Create a new secret. The name of the secret is required. For example:
 
-    Examples:
     ```bash
-    flyte create secret my_secret --value my_value
+    $ 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.
-    Example:
+
+    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
+    $ flyte create secret my_secret --from-file /path/to/secret_file
     ```
-    Secret types can be used to create specific types of secrets. Some secrets are useful for image pull, while some
-    are `regular` / general purpose secrets.
-    Example:
+
+    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
+    $ flyte create secret my_secret --type image_pull
     ```
     """
     from flyte.remote import Secret
@@ -61,28 +65,28 @@ 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 configuration file.",
 )
 @click.option(
     "-o",
     "--output",
     type=click.Path(exists=False, writable=True),
     default=Path.cwd() / "config.yaml",
-    help="Path to the output dir where the config will be saved, defaults to current directory.",
+    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 the config file if it already exists.",
+    help="Force overwrite of the configuration file if it already exists.",
     show_default=True,
-    prompt="Are you sure you want to overwrite the config file?",
+    prompt="Are you sure you want to overwrite the configuration file?",
     confirmation_prompt=True,
 )
 def config(
@@ -95,7 +99,7 @@ def config(
     force: bool = False,
 ):
     """
-    This command creates a configuration file for Flyte CLI.
+    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.
     """
@@ -104,18 +108,29 @@ def config(
     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 = {
-            "admin": {
-                "endpoint": endpoint,
-                "insecure": insecure,
-            },
-            "task": {
-                "org": org,
-                "project": project,
-                "domain": domain,
-            },
-        }
+        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}")

flyte/cli/_delete.py

@@ -6,7 +6,7 @@ import flyte.cli._common as common
 @click.group(name="delete")
 def delete():
     """
-    The delete subcommand allows you to remove resources from a Flyte deployment.
+    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, the name of the secret is required.
+    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