flyte 2.0.0b9__py3-none-any.whl → 2.0.0b14__py3-none-any.whl

flyte/cli/_gen.py

@@ -67,8 +67,13 @@ def markdown(cfg: common.CLIConfig):
     output_verb_groups: dict[str, list[str]] = {}
     output_noun_groups: dict[str, list[str]] = {}
 
+    processed = []
     commands = [*[("flyte", ctx.command)], *walk_commands(ctx)]
     for cmd_path, cmd in commands:
+        if cmd in processed:
+            # We already processed this command, skip it
+            continue
+        processed.append(cmd)
         output.append("")
 
         cmd_path_parts = cmd_path.split(" ")
@@ -136,7 +141,11 @@ def markdown(cfg: common.CLIConfig):
         output_verb_index.append("| ------ | -- |")
         for verb, nouns in output_verb_groups.items():
             entries = [f"[`{noun}`](#flyte-{verb}-{noun})" for noun in nouns]
-            output_verb_index.append(f"| `{verb}` | {', '.join(entries)}  |")
+            if len(entries) == 0:
+                verb_link = f"[`{verb}`](#flyte-{verb})"
+                output_verb_index.append(f"| {verb_link} | - |")
+            else:
+                output_verb_index.append(f"| `{verb}` | {', '.join(entries)}  |")
 
     output_noun_index = []
 

flyte/cli/_get.py

@@ -5,8 +5,6 @@ import rich_click as click
 from rich.console import Console
 from rich.pretty import pretty_repr
 
-import flyte.remote._action
-
 from . import _common as common
 
 
@@ -76,7 +74,7 @@ def run(
     console = Console()
     if name:
         details = RunDetails.get(name=name)
-        console.print(pretty_repr(details))
+        console.print(common.format(f"Run {name}", [details], "json"))
     else:
         console.print(common.format("Runs", Run.listall(limit=limit), cfg.output_format))
 
@@ -110,7 +108,7 @@ def task(
             if v is None:
                 raise click.BadParameter(f"Task {name} not found.")
             t = v.fetch()
-            console.print(pretty_repr(t))
+            console.print(common.format(f"Task {name}", [t], "json"))
         else:
             console.print(common.format("Tasks", Task.listall(by_task_name=name, limit=limit), cfg.output_format))
     else:
@@ -131,19 +129,22 @@ def action(
     """
     Get all actions for a run or details for a specific action.
     """
+    import flyte.remote as remote
 
     cfg.init(project=project, domain=domain)
 
     console = Console()
     if action_name:
-        console.print(pretty_repr(flyte.remote._action.Action.get(run_name=run_name, name=action_name)))
-    else:
-        # List all actions for the run
         console.print(
             common.format(
-                f"Actions for {run_name}", flyte.remote._action.Action.listall(for_run_name=run_name), cfg.output_format
+                f"Action {run_name}.{action_name}", [remote.Action.get(run_name=run_name, name=action_name)], "json"
             )
         )
+    else:
+        # List all actions for the run
+        console.print(
+            common.format(f"Actions for {run_name}", remote.Action.listall(for_run_name=run_name), cfg.output_format)
+        )
 
 
 @get.command(cls=common.CommandBase)
@@ -211,7 +212,7 @@ def logs(
             task.cancel()
 
     if action_name:
-        obj = flyte.remote._action.Action.get(run_name=run_name, name=action_name)
+        obj = remote.Action.get(run_name=run_name, name=action_name)
     else:
         obj = remote.Run.get(run_name)
     asyncio.run(_run_log_view(obj))
@@ -235,7 +236,7 @@ def secret(
 
     console = Console()
     if name:
-        console.print(pretty_repr(remote.Secret.get(name)))
+        console.print(common.format("Secret", [remote.Secret.get(name)], "json"))
     else:
         console.print(common.format("Secrets", remote.Secret.listall(), cfg.output_format))
 
@@ -275,24 +276,25 @@ def io(
         raise click.BadParameter("Cannot use both --inputs-only and --outputs-only")
 
     import flyte.remote as remote
+    from flyte.remote import ActionDetails, ActionInputs, ActionOutputs
 
     cfg.init(project=project, domain=domain)
     console = Console()
     if action_name:
-        obj = flyte.remote._action.ActionDetails.get(run_name=run_name, name=action_name)
+        obj = ActionDetails.get(run_name=run_name, name=action_name)
     else:
         obj = remote.RunDetails.get(run_name)
 
     async def _get_io(
-        details: Union[remote.RunDetails, flyte.remote._action.ActionDetails],
-    ) -> Tuple[flyte.remote._action.ActionInputs | None, flyte.remote._action.ActionOutputs | None | str]:
+        details: Union[remote.RunDetails, ActionDetails],
+    ) -> Tuple[ActionInputs | None, ActionOutputs | None | str]:
         if inputs_only or outputs_only:
             if inputs_only:
                 return await details.inputs(), None
             elif outputs_only:
                 return None, await details.outputs()
         inputs = await details.inputs()
-        outputs: flyte.remote._action.ActionOutputs | None | str = None
+        outputs: ActionOutputs | None | str = None
         try:
             outputs = await details.outputs()
         except Exception:

flyte/cli/_run.py

@@ -3,12 +3,12 @@ from __future__ import annotations
 import asyncio
 import inspect
 from dataclasses import dataclass, field, fields
+from functools import lru_cache
 from pathlib import Path
 from types import ModuleType
 from typing import Any, Dict, List, cast
 
-import click
-from click import Context, Parameter
+import rich_click as click
 from rich.console import Console
 from typing_extensions import get_args
 
@@ -19,6 +19,34 @@ from . import _common as common
 from ._common import CLIConfig
 from ._params import to_click_option
 
+RUN_REMOTE_CMD = "deployed-task"
+
+
+@lru_cache()
+def _initialize_config(ctx: click.Context, project: str, domain: str):
+    obj: CLIConfig | None = ctx.obj
+    if obj is None:
+        import flyte.config
+
+        obj = CLIConfig(flyte.config.auto(), ctx)
+
+    obj.init(project, domain)
+    return obj
+
+
+@lru_cache()
+def _list_tasks(
+    ctx: click.Context,
+    project: str,
+    domain: str,
+    by_task_name: str | None = None,
+    by_task_env: str | None = None,
+) -> list[str]:
+    import flyte.remote
+
+    _initialize_config(ctx, project, domain)
+    return [task.name for task in flyte.remote.Task.listall(by_task_name=by_task_name, by_task_env=by_task_env)]
+
 
 @dataclass
 class RunArguments:
@@ -60,7 +88,7 @@ class RunArguments:
         },
     )
     follow: bool = field(
-        default=True,
+        default=False,
         metadata={
             "click.option": click.Option(
                 ["--follow", "-f"],
@@ -92,23 +120,17 @@ class RunTaskCommand(click.Command):
         kwargs.pop("name", None)
         super().__init__(obj_name, *args, **kwargs)
 
-    def invoke(self, ctx: Context):
-        obj: CLIConfig = ctx.obj
-        if obj is None:
-            import flyte.config
-
-            obj = CLIConfig(flyte.config.auto(), ctx)
-
-        obj.init(self.run_args.project, self.run_args.domain)
+    def invoke(self, ctx: click.Context):
+        obj: CLIConfig = _initialize_config(ctx, self.run_args.project, self.run_args.domain)
 
         async def _run():
             import flyte
 
-            r = flyte.with_runcontext(
+            r = await flyte.with_runcontext(
                 copy_style=self.run_args.copy_style,
                 mode="local" if self.run_args.local else "remote",
                 name=self.run_args.name,
-            ).run(self.obj, **ctx.params)
+            ).run.aio(self.obj, **ctx.params)
             if isinstance(r, Run) and r.action is not None:
                 console = Console()
                 console.print(
@@ -129,7 +151,7 @@ class RunTaskCommand(click.Command):
 
         asyncio.run(_run())
 
-    def get_params(self, ctx: Context) -> List[Parameter]:
+    def get_params(self, ctx: click.Context) -> List[click.Parameter]:
         # Note this function may be called multiple times by click.
         task = self.obj
         from .._internal.runtime.types_serde import transform_native_to_typed_interface
@@ -139,7 +161,7 @@ class RunTaskCommand(click.Command):
             return super().get_params(ctx)
         inputs_interface = task.native_interface.inputs
 
-        params: List[Parameter] = []
+        params: List[click.Parameter] = []
         for name, var in interface.inputs.variables.items():
             default_val = None
             if inputs_interface[name][1] is not inspect._empty:
@@ -156,8 +178,9 @@ class TaskPerFileGroup(common.ObjectsPerFileGroup):
     """
 
     def __init__(self, filename: Path, run_args: RunArguments, *args, **kwargs):
-        args = (filename, *args)
-        super().__init__(*args, **kwargs)
+        if filename.is_absolute():
+            filename = filename.relative_to(Path.cwd())
+        super().__init__(*(filename, *args), **kwargs)
         self.run_args = run_args
 
     def _filter_objects(self, module: ModuleType) -> Dict[str, Any]:
@@ -173,6 +196,172 @@ class TaskPerFileGroup(common.ObjectsPerFileGroup):
         )
 
 
+class RunReferenceTaskCommand(click.Command):
+    def __init__(self, task_name: str, run_args: RunArguments, version: str | None, *args, **kwargs):
+        self.task_name = task_name
+        self.run_args = run_args
+        self.version = version
+
+        super().__init__(*args, **kwargs)
+
+    def invoke(self, ctx: click.Context):
+        obj: CLIConfig = _initialize_config(ctx, self.run_args.project, self.run_args.domain)
+
+        async def _run():
+            import flyte
+            import flyte.remote
+
+            task = flyte.remote.Task.get(self.task_name, version=self.version, auto_version="latest")
+
+            r = await flyte.with_runcontext(
+                copy_style=self.run_args.copy_style,
+                mode="local" if self.run_args.local else "remote",
+                name=self.run_args.name,
+            ).run.aio(task, **ctx.params)
+            if isinstance(r, Run) and r.action is not None:
+                console = Console()
+                console.print(
+                    common.get_panel(
+                        "Run",
+                        f"[green bold]Created Run: {r.name} [/green bold] "
+                        f"(Project: {r.action.action_id.run.project}, Domain: {r.action.action_id.run.domain})\n"
+                        f"➡️  [blue bold][link={r.url}]{r.url}[/link][/blue bold]",
+                        obj.output_format,
+                    )
+                )
+                if self.run_args.follow:
+                    console.print(
+                        "[dim]Log streaming enabled, will wait for task to start running "
+                        "and log stream to be available[/dim]"
+                    )
+                    await r.show_logs.aio(max_lines=30, show_ts=True, raw=False)
+
+        asyncio.run(_run())
+
+    def get_params(self, ctx: click.Context) -> List[click.Parameter]:
+        # Note this function may be called multiple times by click.
+        import flyte.remote
+        from flyte._internal.runtime.types_serde import transform_native_to_typed_interface
+
+        _initialize_config(ctx, self.run_args.project, self.run_args.domain)
+
+        task = flyte.remote.Task.get(self.task_name, auto_version="latest")
+        task_details = task.fetch()
+
+        interface = transform_native_to_typed_interface(task_details.interface)
+        if interface is None:
+            return super().get_params(ctx)
+        inputs_interface = task_details.interface.inputs
+
+        params: List[click.Parameter] = []
+        for name, var in interface.inputs.variables.items():
+            default_val = None
+            if inputs_interface[name][1] is not inspect._empty:
+                default_val = inputs_interface[name][1]
+            params.append(to_click_option(name, var, inputs_interface[name][0], default_val))
+
+        self.params = params
+        return super().get_params(ctx)
+
+
+class ReferenceEnvGroup(common.GroupBase):
+    def __init__(self, name: str, *args, run_args, env: str, **kwargs):
+        super().__init__(*args, **kwargs)
+        self.name = name
+        self.env = env
+        self.run_args = run_args
+
+    def list_commands(self, ctx):
+        return _list_tasks(ctx, self.run_args.project, self.run_args.domain, by_task_env=self.env)
+
+    def get_command(self, ctx, name):
+        return RunReferenceTaskCommand(
+            task_name=name,
+            run_args=self.run_args,
+            name=name,
+            version=None,
+            help=f"Run deployed task '{name}' from the Flyte backend",
+        )
+
+
+class ReferenceTaskGroup(common.GroupBase):
+    """
+    Group that creates a command for each reference task in the current directory that is not __init__.py.
+    """
+
+    def __init__(self, name: str, *args, run_args, tasks: list[str] | None = None, **kwargs):
+        super().__init__(*args, **kwargs)
+        self.name = name
+        self.run_args = run_args
+
+    def list_commands(self, ctx):
+        # list envs of all reference tasks
+        envs = []
+        for task in _list_tasks(ctx, self.run_args.project, self.run_args.domain):
+            env = task.split(".")[0]
+            if env not in envs:
+                envs.append(env)
+        return envs
+
+    @staticmethod
+    def _parse_task_name(task_name: str) -> tuple[str, str | None, str | None]:
+        import re
+
+        pattern = r"^([^.:]+)(?:\.([^:]+))?(?::(.+))?$"
+        match = re.match(pattern, task_name)
+        if not match:
+            raise click.BadParameter(f"Invalid task name format: {task_name}")
+        return match.group(1), match.group(2), match.group(3)
+
+    def _env_is_task(self, ctx: click.Context, env: str) -> bool:
+        # check if the env name is the full task name, since sometimes task
+        # names don't have an environment prefix
+        tasks = [*_list_tasks(ctx, self.run_args.project, self.run_args.domain, by_task_name=env)]
+        return len(tasks) > 0
+
+    def get_command(self, ctx, name):
+        env, task, version = self._parse_task_name(name)
+        match env, task, version:
+            case env, None, None:
+                if self._env_is_task(ctx, env):
+                    # this handles cases where task names do not have a environment prefix
+                    task_name = env
+                    return RunReferenceTaskCommand(
+                        task_name=task_name,
+                        run_args=self.run_args,
+                        name=task_name,
+                        version=None,
+                        help=f"Run reference task `{task_name}` from the Flyte backend",
+                    )
+                else:
+                    return ReferenceEnvGroup(
+                        name=name,
+                        run_args=self.run_args,
+                        env=env,
+                        help=f"Run reference tasks in the `{env}` environment from the Flyte backend",
+                    )
+            case env, task, None:
+                task_name = f"{env}.{task}"
+                return RunReferenceTaskCommand(
+                    task_name=task_name,
+                    run_args=self.run_args,
+                    name=task_name,
+                    version=None,
+                    help=f"Run reference task '{task_name}' from the Flyte backend",
+                )
+            case env, task, version:
+                task_name = f"{env}.{task}"
+                return RunReferenceTaskCommand(
+                    task_name=task_name,
+                    run_args=self.run_args,
+                    version=version,
+                    name=f"{task_name}:{version}",
+                    help=f"Run reference task '{task_name}' from the Flyte backend",
+                )
+            case _:
+                raise click.BadParameter(f"Invalid task name format: {task_name}")
+
+
 class TaskFiles(common.FileGroup):
     """
     Group that creates a command for each file in the current directory that is not __init__.py.
@@ -191,25 +380,45 @@ class TaskFiles(common.FileGroup):
         kwargs["params"].extend(RunArguments.options())
         super().__init__(*args, directory=directory, **kwargs)
 
-    def get_command(self, ctx, filename):
+    def list_commands(self, ctx):
+        v = [
+            RUN_REMOTE_CMD,
+            *super().list_commands(ctx),
+        ]
+        return v
+
+    def get_command(self, ctx, cmd_name):
         run_args = RunArguments.from_dict(ctx.params)
-        fp = Path(filename)
+
+        if cmd_name == RUN_REMOTE_CMD:
+            return ReferenceTaskGroup(
+                name=cmd_name,
+                run_args=run_args,
+                help="Run reference task from the Flyte backend",
+            )
+
+        fp = Path(cmd_name)
         if not fp.exists():
-            raise click.BadParameter(f"File {filename} does not exist")
+            raise click.BadParameter(f"File {cmd_name} does not exist")
         if fp.is_dir():
-            return TaskFiles(directory=fp)
+            return TaskFiles(
+                directory=fp,
+                help=f"Run `*.py` file inside the {fp} directory",
+            )
         return TaskPerFileGroup(
             filename=fp,
             run_args=run_args,
-            name=filename,
-            help=f"Run, functions decorated with `env.task` in {filename}",
+            name=cmd_name,
+            help=f"Run functions decorated with `env.task` in {cmd_name}",
         )
 
 
 run = TaskFiles(
     name="run",
-    help="""
-Run a task from a python file.
+    help=f"""
+Run a task from a python file or deployed task.
+
+To run a remote task that already exists in Flyte, use the {RUN_REMOTE_CMD} command:
 
 Example usage:
 
@@ -227,6 +436,30 @@ Flyte environment:
 flyte run --local hello.py my_task --arg1 value1 --arg2 value2
 ```
 
+To run tasks that you've already deployed to Flyte, use the {RUN_REMOTE_CMD} command:
+
+```bash
+flyte run {RUN_REMOTE_CMD} my_env.my_task --arg1 value1 --arg2 value2
+```
+
+To run a specific version of a deployed task, use the `env.task:version` syntax:
+
+```bash
+flyte run {RUN_REMOTE_CMD} my_env.my_task:xyz123 --arg1 value1 --arg2 value2
+```
+
+You can specify the `--config` flag to point to a specific Flyte cluster:
+
+```bash
+flyte run --config my-config.yaml {RUN_REMOTE_CMD} ...
+```
+
+You can discover what deployed tasks are available by running:
+
+```bash
+flyte run {RUN_REMOTE_CMD}
+```
+
 Other arguments to the run command are listed below.
 
 Arguments for the task itself are provided after the task name and can be retrieved using `--help`. For example:

flyte/models.py

@@ -166,6 +166,7 @@ class TaskContext:
     action: ActionID
     version: str
     raw_data_path: RawDataPath
+    input_path: str | None = None
     output_path: str
     run_base_dir: str
     report: Report
@@ -175,6 +176,7 @@ class TaskContext:
     compiled_image_cache: ImageCache | None = None
     data: Dict[str, Any] = field(default_factory=dict)
     mode: Literal["local", "remote", "hybrid"] = "remote"
+    interactive_mode: bool = False
 
     def replace(self, **kwargs) -> TaskContext:
         if "data" in kwargs:
@@ -191,6 +193,13 @@ class TaskContext:
     def __getitem__(self, key: str) -> Optional[Any]:
         return self.data.get(key)
 
+    def is_in_cluster(self):
+        """
+        Check if the task is running in a cluster.
+        :return: bool
+        """
+        return self.mode == "remote"
+
 
 @rich.repr.auto
 @dataclass(frozen=True, kw_only=True)

flyte/remote/_client/auth/_authenticators/base.py

@@ -34,6 +34,7 @@ class Authenticator(object):
         http_proxy_url: typing.Optional[str] = None,
         verify: bool = True,
         ca_cert_path: typing.Optional[str] = None,
+        default_header_key: str = "authorization",
         **kwargs,
     ):
         """
@@ -80,6 +81,7 @@ class Authenticator(object):
         self._http_session = http_session or get_async_session(**kwargs)
         # Id for tracking credential refresh state
         self._creds_id = self._creds.id if self._creds else None
+        self._default_header_key = default_header_key
 
     async def _resolve_config(self) -> ClientConfig:
         """
@@ -131,10 +133,14 @@ class Authenticator(object):
         """
         creds = self.get_credentials()
         if creds:
-            cfg = await self._resolve_config()
+            header_key = self._default_header_key
+            if self._resolved_config is not None:
+                # We only resolve the config during authentication flow, to avoid unnecessary network calls
+                # and usually the header_key is consistent.
+                header_key = self._resolved_config.header_key
             return GrpcAuthMetadata(
                 creds_id=creds.id,
-                pairs=Metadata((cfg.header_key, f"Bearer {creds.access_token}")),
+                pairs=Metadata((header_key, f"Bearer {creds.access_token}")),
             )
         return None
 

flyte/remote/_client/auth/_authenticators/device_code.py

@@ -81,7 +81,7 @@ class DeviceCodeAuthenticator(Authenticator):
                     for_endpoint=self._endpoint,
                 )
             except (AuthenticationError, AuthenticationPending):
-                logger.warning("Failed to refresh token. Kicking off a full authorization flow.")
+                logger.warning("Logging in...")
 
         """Fall back to device flow"""
         resp = await token_client.get_device_code(

flyte/remote/_client/auth/_authenticators/pkce.py

@@ -123,7 +123,7 @@ class PKCEAuthenticator(Authenticator):
             try:
                 return await self._auth_client.refresh_access_token(self._creds)
             except AccessTokenNotFoundError:
-                logger.warning("Failed to refresh token. Kicking off a full authorization flow.")
+                logger.warning("Logging in...")
 
         return await self._auth_client.get_creds_from_remote()
 

flyte/remote/_client/auth/_channel.py

@@ -1,4 +1,3 @@
-import os
 import ssl
 import typing
 
@@ -16,11 +15,6 @@ from ._authenticators.factory import (
     get_async_proxy_authenticator,
 )
 
-# Set environment variables for gRPC, this reduces log spew and avoids unnecessary warnings
-if "GRPC_VERBOSITY" not in os.environ:
-    os.environ["GRPC_VERBOSITY"] = "ERROR"
-    os.environ["GRPC_CPP_MIN_LOG_LEVEL"] = "ERROR"
-
 # Initialize gRPC AIO early enough so it can be used in the main thread
 init_grpc_aio()
 

flyte/remote/_client/auth/_client_config.py

@@ -1,3 +1,4 @@
+import asyncio
 import typing
 from abc import abstractmethod
 
@@ -69,8 +70,9 @@ class RemoteClientConfigStore(ClientConfigStore):
         Retrieves the ClientConfig from the given grpc.Channel assuming  AuthMetadataService is available
         """
         metadata_service = AuthMetadataServiceStub(self._unauthenticated_channel)
-        public_client_config = await metadata_service.GetPublicClientConfig(PublicClientAuthConfigRequest())
-        oauth2_metadata = await metadata_service.GetOAuth2Metadata(OAuth2MetadataRequest())
+        oauth2_metadata_task = metadata_service.GetOAuth2Metadata(OAuth2MetadataRequest())
+        public_client_config_task = metadata_service.GetPublicClientConfig(PublicClientAuthConfigRequest())
+        oauth2_metadata, public_client_config = await asyncio.gather(oauth2_metadata_task, public_client_config_task)
         return ClientConfig(
             token_endpoint=oauth2_metadata.token_endpoint,
             authorization_endpoint=oauth2_metadata.authorization_endpoint,

flyte/remote/_client/controlplane.py

@@ -1,5 +1,19 @@
 from __future__ import annotations
 
+import os
+
+# Set environment variables for gRPC, this reduces log spew and avoids unnecessary warnings
+# before importing grpc
+if "GRPC_VERBOSITY" not in os.environ:
+    os.environ["GRPC_VERBOSITY"] = "ERROR"
+    os.environ["GRPC_CPP_MIN_LOG_LEVEL"] = "ERROR"
+    # Disable fork support (stops "skipping fork() handlers")
+    os.environ["GRPC_ENABLE_FORK_SUPPORT"] = "0"
+    # Reduce absl/glog verbosity
+    os.environ["GLOG_minloglevel"] = "2"
+    os.environ["ABSL_LOG"] = "0"
+#### Has to be before grpc
+
 import grpc
 from flyteidl.service import admin_pb2_grpc, dataproxy_pb2_grpc