flyte 2.0.0b13__py3-none-any.whl → 2.0.0b30__py3-none-any.whl

flyte/cli/_params.py

@@ -15,9 +15,9 @@ from typing import get_args
 import rich_click as click
 import yaml
 from click import Parameter
-from flyteidl.core.interface_pb2 import Variable
-from flyteidl.core.literals_pb2 import Literal
-from flyteidl.core.types_pb2 import BlobType, LiteralType, SimpleType
+from flyteidl2.core.interface_pb2 import Variable
+from flyteidl2.core.literals_pb2 import Literal
+from flyteidl2.core.types_pb2 import BlobType, LiteralType, SimpleType
 from google.protobuf.json_format import MessageToDict
 from mashumaro.codecs.json import JSONEncoder
 
@@ -283,13 +283,17 @@ class UnionParamType(click.ParamType):
     A composite type that allows for multiple types to be specified. This is used for union types.
     """
 
-    def __init__(self, types: typing.List[click.ParamType]):
+    def __init__(self, types: typing.List[click.ParamType | None]):
         super().__init__()
         self._types = self._sort_precedence(types)
-        self.name = "|".join([t.name for t in self._types])
+        self.name = "|".join([t.name for t in self._types if t is not None])
+        self.optional = False
+        if None in types:
+            self.name = f"Optional[{self.name}]"
+            self.optional = True
 
     @staticmethod
-    def _sort_precedence(tp: typing.List[click.ParamType]) -> typing.List[click.ParamType]:
+    def _sort_precedence(tp: typing.List[click.ParamType | None]) -> typing.List[click.ParamType]:
         unprocessed = []
         str_types = []
         others = []
@@ -311,6 +315,8 @@ class UnionParamType(click.ParamType):
         """
         for p in self._types:
             try:
+                if p is None and value is None:
+                    return None
                 return p.convert(value, param, ctx)
             except Exception as e:
                 logger.debug(f"Ignoring conversion error for type {p} trying other variants in Union. Error: {e}")
@@ -433,7 +439,10 @@ def literal_type_to_click_type(lt: LiteralType, python_type: typing.Type) -> cli
         for i in range(len(lt.union_type.variants)):
             variant = lt.union_type.variants[i]
             variant_python_type = typing.get_args(python_type)[i]
-            cts.append(literal_type_to_click_type(variant, variant_python_type))
+            if variant_python_type is type(None):
+                cts.append(None)
+            else:
+                cts.append(literal_type_to_click_type(variant, variant_python_type))
         return UnionParamType(cts)
 
     if lt.HasField("enum_type"):
@@ -461,6 +470,9 @@ class FlyteLiteralConverter(object):
     def is_bool(self) -> bool:
         return self.click_type == click.BOOL
 
+    def is_optional(self) -> bool:
+        return isinstance(self.click_type, UnionParamType) and self.click_type.optional
+
     def convert(
         self, ctx: click.Context, param: typing.Optional[click.Parameter], value: typing.Any
     ) -> typing.Union[Literal, typing.Any]:
@@ -493,7 +505,7 @@ def to_click_option(
     This handles converting workflow input types to supported click parameters with callbacks to initialize
     the input values to their expected types.
     """
-    from flyteidl.core.types_pb2 import SimpleType
+    from flyteidl2.core.types_pb2 import SimpleType
 
     if input_name != input_name.lower():
         # Click does not support uppercase option names: https://github.com/pallets/click/issues/837
@@ -519,15 +531,19 @@ def to_click_option(
         if literal_var.type.metadata:
             description_extra = f": {MessageToDict(literal_var.type.metadata)}"
 
-    # If a query has been specified, the input is never strictly required at this layer
     required = False if default_val is not None else True
     is_flag: typing.Optional[bool] = None
+    param_decls = [f"--{input_name}"]
     if literal_converter.is_bool():
         required = False
         is_flag = True
+        if default_val is True:
+            param_decls = [f"--{input_name}/--no-{input_name}"]
+    if literal_converter.is_optional():
+        required = False
 
     return click.Option(
-        param_decls=[f"--{input_name}"],
+        param_decls=param_decls,
         type=literal_converter.click_type,
         is_flag=is_flag,
         default=default_val,

flyte/cli/_plugins.py

@@ -0,0 +1,209 @@
+"""CLI Plugin System for Flyte.
+
+This module provides a plugin system that allows external packages to:
+1. Register new top-level CLI commands (e.g., flyte my-command)
+2. Register new subcommands in existing groups (e.g., flyte get my-object)
+3. Modify behavior of existing commands via hooks
+
+Plugins are discovered via Python entry points.
+
+Entry Point Groups:
+- flyte.plugins.cli.commands: Register new commands
+  - Entry point name "foo" -> flyte foo (top-level command)
+  - Entry point name "get.bar" -> flyte get bar (adds subcommand to get group)
+  - Note: At most one dot is supported. For nested groups, register the entire
+    group hierarchy as a top-level command (without dots).
+
+- flyte.plugins.cli.hooks: Modify existing commands
+  - Entry point name "run" -> modifies flyte run
+  - Entry point name "get.project" -> modifies flyte get project
+  - Note: At most one dot is supported.
+
+Example Plugin Package:
+    # In your-plugin/pyproject.toml
+    [project.entry-points."flyte.plugins.cli.commands"]
+    my-command = "your_plugin.cli:my_command"
+    get.my-object = "your_plugin.cli:get_my_object"
+
+    [project.entry-points."flyte.plugins.cli.hooks"]
+    run = "your_plugin.hooks:modify_run"
+
+    # In your-plugin/your_plugin/cli.py
+    import rich_click as click
+
+    @click.command()
+    def my_command():
+        '''My custom top-level command.'''
+        click.echo("Hello from plugin!")
+
+    @click.command()
+    def get_my_object():
+        '''Get my custom object.'''
+        click.echo("Getting my object...")
+
+    # In your-plugin/your_plugin/hooks.py
+    def modify_run(command):
+        '''Add behavior to flyte run command.'''
+        # Wrap invoke() instead of callback to ensure Click's full machinery runs
+        original_invoke = command.invoke
+
+        def wrapper(ctx):
+            # Do something before
+            click.echo("Plugin: Starting task...")
+
+            result = original_invoke(ctx)
+
+            # Do something after
+            click.echo("Plugin: Task completed!")
+            return result
+
+        command.invoke = wrapper
+        return command
+"""
+
+from importlib.metadata import entry_points
+from typing import Callable
+
+import rich_click as click
+
+from flyte._logging import logger
+
+# Type alias for command hooks
+CommandHook = Callable[[click.Command], click.Command]
+
+
+def discover_and_register_plugins(root_group: click.Group):
+    """
+    Discover all CLI plugins from installed packages and register them.
+
+    This function:
+    1. Discovers command plugins and adds them to the CLI
+    2. Discovers hook plugins and applies them to existing commands
+
+    Args:
+        root_group: The root Click command group (main CLI group)
+    """
+    _load_command_plugins(root_group)
+    _load_hook_plugins(root_group)
+
+
+def _load_command_plugins(root_group: click.Group):
+    """Load and register command plugins."""
+    for ep in entry_points(group="flyte.plugins.cli.commands"):
+        try:
+            command = ep.load()
+            if not isinstance(command, click.Command):
+                logger.warning(f"Plugin {ep.name} did not return a click.Command, got {type(command).__name__}")
+                continue
+
+            # Check if this is a subcommand (contains dot notation)
+            if "." in ep.name:
+                group_name, command_name = ep.name.split(".", 1)
+
+                # Validate: only support one level of nesting (group.command)
+                if "." in command_name:
+                    logger.error(
+                        f"Plugin {ep.name} uses multiple dots, which is not supported. "
+                        f"Use at most one dot (e.g., 'group.command'). "
+                        f"For nested groups, register the entire group hierarchy as a top-level command."
+                    )
+                    continue
+
+                _add_subcommand_to_group(root_group, group_name, command_name, command)
+            else:
+                # Top-level command
+                root_group.add_command(command, name=ep.name)
+                logger.info(f"Registered plugin command: flyte {ep.name}")
+
+        except Exception as e:
+            logger.error(f"Failed to load plugin command {ep.name}: {e}")
+
+
+def _load_hook_plugins(root_group: click.Group):
+    """Load and apply hook plugins to existing commands."""
+    for ep in entry_points(group="flyte.plugins.cli.hooks"):
+        try:
+            hook = ep.load()
+            if not callable(hook):
+                logger.warning(f"Plugin hook {ep.name} is not callable")
+                continue
+
+            # Check if this is a subcommand hook (contains dot notation)
+            if "." in ep.name:
+                group_name, command_name = ep.name.split(".", 1)
+
+                # Validate: only support one level of nesting (group.command)
+                if "." in command_name:
+                    logger.error(
+                        f"Plugin hook {ep.name} uses multiple dots, which is not supported. "
+                        f"Use at most one dot (e.g., 'group.command')."
+                    )
+                    continue
+
+                _apply_hook_to_subcommand(root_group, group_name, command_name, hook)
+            else:
+                # Top-level command hook
+                _apply_hook_to_command(root_group, ep.name, hook)
+
+        except Exception as e:
+            logger.error(f"Failed to apply hook {ep.name}: {e}")
+
+
+def _add_subcommand_to_group(root_group: click.Group, group_name: str, command_name: str, command: click.Command):
+    """Add a subcommand to an existing command group."""
+    if group_name not in root_group.commands:
+        logger.warning(f"Cannot add plugin subcommand '{command_name}' - group '{group_name}' does not exist")
+        return
+
+    group = root_group.commands[group_name]
+    if not isinstance(group, click.Group):
+        logger.warning(f"Cannot add plugin subcommand '{command_name}' - '{group_name}' is not a command group")
+        return
+
+    group.add_command(command, name=command_name)
+    # lower to debug later
+    logger.info(f"Registered plugin subcommand: flyte {group_name} {command_name}")
+
+
+def _apply_hook_to_command(root_group: click.Group, command_name: str, hook: CommandHook):
+    """Apply a hook to a top-level command."""
+    if command_name not in root_group.commands:
+        logger.warning(f"Cannot apply hook - command '{command_name}' does not exist")
+        return
+
+    original_command = root_group.commands[command_name]
+    try:
+        modified_command = hook(original_command)
+        root_group.commands[command_name] = modified_command
+        # lower to debug later
+        logger.info(f"Applied hook to command: flyte {command_name}")
+    except Exception as e:
+        logger.error(f"Hook failed for command {command_name}: {e}")
+        root_group.commands[command_name] = original_command
+
+
+def _apply_hook_to_subcommand(root_group: click.Group, group_name: str, command_name: str, hook: CommandHook):
+    """Apply a hook to a subcommand within a group."""
+    if group_name not in root_group.commands:
+        logger.warning(f"Cannot apply hook - group '{group_name}' does not exist")
+        return
+
+    group = root_group.commands[group_name]
+    if not isinstance(group, click.Group):
+        logger.warning(f"Cannot apply hook - '{group_name}' is not a command group")
+        return
+
+    if command_name not in group.commands:
+        logger.warning(f"Cannot apply hook - subcommand '{command_name}' does not exist in group '{group_name}'")
+        return
+
+    original_command = group.commands[command_name]
+    if original_command.callback is not None:
+        original_command.callback()
+    try:
+        modified_command = hook(original_command)
+        group.commands[command_name] = modified_command
+        logger.info(f"Applied hook to subcommand: flyte {group_name} {command_name}")
+    except Exception as e:
+        logger.error(f"Hook failed for subcommand {group_name}.{command_name}: {e}")
+        group.commands[command_name] = original_command

flyte/cli/_run.py

@@ -8,36 +8,34 @@ from pathlib import Path
 from types import ModuleType
 from typing import Any, Dict, List, cast
 
-import click
-from click import Context, Parameter
-from rich.console import Console
+import rich_click as click
 from typing_extensions import get_args
 
 from .._code_bundle._utils import CopyFiles
 from .._task import TaskTemplate
 from ..remote import Run
 from . import _common as common
-from ._common import CLIConfig
+from ._common import CLIConfig, initialize_config
 from ._params import to_click_option
 
 RUN_REMOTE_CMD = "deployed-task"
 
 
 @lru_cache()
-def _initialize_config(ctx: Context, project: str, domain: str):
+def _initialize_config(ctx: click.Context, project: str, domain: str, root_dir: str | None = None):
     obj: CLIConfig | None = ctx.obj
     if obj is None:
         import flyte.config
 
         obj = CLIConfig(flyte.config.auto(), ctx)
 
-    obj.init(project, domain)
+    obj.init(project, domain, root_dir)
     return obj
 
 
 @lru_cache()
 def _list_tasks(
-    ctx: Context,
+    ctx: click.Context,
     project: str,
     domain: str,
     by_task_name: str | None = None,
@@ -45,7 +43,7 @@ def _list_tasks(
 ) -> list[str]:
     import flyte.remote
 
-    _initialize_config(ctx, project, domain)
+    common.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)]
 
 
@@ -78,6 +76,36 @@ class RunArguments:
             )
         },
     )
+    root_dir: str | None = field(
+        default=None,
+        metadata={
+            "click.option": click.Option(
+                ["--root-dir"],
+                type=str,
+                help="Override the root source directory, helpful when working with monorepos.",
+            )
+        },
+    )
+    raw_data_path: str | None = field(
+        default=None,
+        metadata={
+            "click.option": click.Option(
+                ["--raw-data-path"],
+                type=str,
+                help="Override the output prefix used to store offloaded data types. e.g. s3://bucket/",
+            )
+        },
+    )
+    service_account: str | None = field(
+        default=None,
+        metadata={
+            "click.option": click.Option(
+                ["--service-account"],
+                type=str,
+                help="Kubernetes service account. If not provided, the configured default will be used",
+            )
+        },
+    )
     name: str | None = field(
         default=None,
         metadata={
@@ -100,10 +128,35 @@ class RunArguments:
             )
         },
     )
+    image: List[str] = field(
+        default_factory=list,
+        metadata={
+            "click.option": click.Option(
+                ["--image"],
+                type=str,
+                multiple=True,
+                help="Image to be used in the run. Format: imagename=imageuri. Can be specified multiple times.",
+            )
+        },
+    )
+    no_sync_local_sys_paths: bool = field(
+        default=True,
+        metadata={
+            "click.option": click.Option(
+                ["--no-sync-local-sys-paths"],
+                is_flag=True,
+                flag_value=True,
+                default=False,
+                help="Disable synchronization of local sys.path entries under the root directory "
+                "to the remote container.",
+            )
+        },
+    )
 
     @classmethod
     def from_dict(cls, d: Dict[str, Any]) -> RunArguments:
-        return cls(**d)
+        modified = {k: v for k, v in d.items() if k in {f.name for f in fields(cls)}}
+        return cls(**modified)
 
     @classmethod
     def options(cls) -> List[click.Option]:
@@ -113,7 +166,7 @@ class RunArguments:
         return [common.get_option_from_metadata(f.metadata) for f in fields(cls) if f.metadata]
 
 
-class RunTaskCommand(click.Command):
+class RunTaskCommand(click.RichCommand):
     def __init__(self, obj_name: str, obj: Any, run_args: RunArguments, *args, **kwargs):
         self.obj_name = obj_name
         self.obj = cast(TaskTemplate, obj)
@@ -121,19 +174,39 @@ class RunTaskCommand(click.Command):
         kwargs.pop("name", None)
         super().__init__(obj_name, *args, **kwargs)
 
-    def invoke(self, ctx: Context):
-        obj: CLIConfig = _initialize_config(ctx, 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,
+            self.run_args.root_dir,
+            tuple(self.run_args.image) or None,
+            not self.run_args.no_sync_local_sys_paths,
+        )
 
         async def _run():
             import flyte
 
+            console = common.get_console()
             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,
+                raw_data_path=self.run_args.raw_data_path,
+                service_account=self.run_args.service_account,
+                log_format=obj.log_format,
             ).run.aio(self.obj, **ctx.params)
+            if self.run_args.local:
+                console.print(
+                    common.get_panel(
+                        "Local Run",
+                        f"[green]Completed Local Run, data stored in path: {r.url} [/green] \n"
+                        f"➡️  Outputs: {r.outputs()}",
+                        obj.output_format,
+                    )
+                )
+                return
             if isinstance(r, Run) and r.action is not None:
-                console = Console()
                 console.print(
                     common.get_panel(
                         "Run",
@@ -152,7 +225,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
@@ -162,7 +235,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:
@@ -187,6 +260,26 @@ class TaskPerFileGroup(common.ObjectsPerFileGroup):
     def _filter_objects(self, module: ModuleType) -> Dict[str, Any]:
         return {k: v for k, v in module.__dict__.items() if isinstance(v, TaskTemplate)}
 
+    def list_commands(self, ctx):
+        common.initialize_config(
+            ctx,
+            self.run_args.project,
+            self.run_args.domain,
+            self.run_args.root_dir,
+            sync_local_sys_paths=not self.run_args.no_sync_local_sys_paths,
+        )
+        return super().list_commands(ctx)
+
+    def get_command(self, ctx, obj_name):
+        common.initialize_config(
+            ctx,
+            self.run_args.project,
+            self.run_args.domain,
+            self.run_args.root_dir,
+            sync_local_sys_paths=not self.run_args.no_sync_local_sys_paths,
+        )
+        return super().get_command(ctx, obj_name)
+
     def _get_command_for_obj(self, ctx: click.Context, obj_name: str, obj: Any) -> click.Command:
         obj = cast(TaskTemplate, obj)
         return RunTaskCommand(
@@ -197,7 +290,7 @@ class TaskPerFileGroup(common.ObjectsPerFileGroup):
         )
 
 
-class RunReferenceTaskCommand(click.Command):
+class RunReferenceTaskCommand(click.RichCommand):
     def __init__(self, task_name: str, run_args: RunArguments, version: str | None, *args, **kwargs):
         self.task_name = task_name
         self.run_args = run_args
@@ -206,13 +299,20 @@ class RunReferenceTaskCommand(click.Command):
         super().__init__(*args, **kwargs)
 
     def invoke(self, ctx: click.Context):
-        obj: CLIConfig = _initialize_config(ctx, self.run_args.project, self.run_args.domain)
+        obj: CLIConfig = common.initialize_config(
+            ctx,
+            self.run_args.project,
+            self.run_args.domain,
+            self.run_args.root_dir,
+            tuple(self.run_args.image) or None,
+            not self.run_args.no_sync_local_sys_paths,
+        )
 
         async def _run():
-            import flyte
             import flyte.remote
 
             task = flyte.remote.Task.get(self.task_name, version=self.version, auto_version="latest")
+            console = common.get_console()
 
             r = await flyte.with_runcontext(
                 copy_style=self.run_args.copy_style,
@@ -220,7 +320,6 @@ class RunReferenceTaskCommand(click.Command):
                 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",
@@ -239,12 +338,17 @@ class RunReferenceTaskCommand(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.
         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)
+        common.initialize_config(
+            ctx,
+            self.run_args.project,
+            self.run_args.domain,
+            sync_local_sys_paths=not self.run_args.no_sync_local_sys_paths,
+        )
 
         task = flyte.remote.Task.get(self.task_name, auto_version="latest")
         task_details = task.fetch()
@@ -254,7 +358,7 @@ class RunReferenceTaskCommand(click.Command):
             return super().get_params(ctx)
         inputs_interface = task_details.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:
@@ -322,7 +426,6 @@ class ReferenceTaskGroup(common.GroupBase):
 
     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):
@@ -383,14 +486,14 @@ class TaskFiles(common.FileGroup):
         super().__init__(*args, directory=directory, **kwargs)
 
     def list_commands(self, ctx):
-        return [
+        v = [
             RUN_REMOTE_CMD,
-            *self.files,
+            *super().list_commands(ctx),
         ]
+        return v
 
     def get_command(self, ctx, cmd_name):
         run_args = RunArguments.from_dict(ctx.params)
-
         if cmd_name == RUN_REMOTE_CMD:
             return ReferenceTaskGroup(
                 name=cmd_name,
@@ -437,6 +540,28 @@ Flyte environment:
 flyte run --local hello.py my_task --arg1 value1 --arg2 value2
 ```
 
+You can provide image mappings with `--image` flag. This allows you to specify
+the image URI for the task environment during CLI execution without changing
+the code. Any images defined with `Image.from_ref_name("name")` will resolve to the
+corresponding URIs you specify here.
+
+```bash
+flyte run hello.py my_task --image my_image=ghcr.io/myorg/my-image:v1.0
+```
+
+If the image name is not provided, it is regarded as a default image and will
+be used when no image is specified in TaskEnvironment:
+
+```bash
+flyte run hello.py my_task --image ghcr.io/myorg/default-image:latest
+```
+
+You can specify multiple image arguments:
+
+```bash
+flyte run hello.py my_task --image ghcr.io/org/default:latest --image gpu=ghcr.io/org/gpu:v2.0
+```
+
 To run tasks that you've already deployed to Flyte, use the {RUN_REMOTE_CMD} command:
 
 ```bash