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

flyte/cli/_run.py

@@ -9,14 +9,13 @@ from types import ModuleType
 from typing import Any, Dict, List, cast
 
 import rich_click as click
-from rich.console import Console
 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"
@@ -44,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)]
 
 
@@ -87,6 +86,26 @@ class RunArguments:
             )
         },
     )
+    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={
@@ -109,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]:
@@ -131,18 +175,38 @@ class RunTaskCommand(click.RichCommand):
         super().__init__(obj_name, *args, **kwargs)
 
     def invoke(self, ctx: click.Context):
-        obj: CLIConfig = _initialize_config(ctx, self.run_args.project, self.run_args.domain, self.run_args.root_dir)
+        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",
@@ -196,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(
@@ -215,13 +299,20 @@ class RunReferenceTaskCommand(click.RichCommand):
         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,
@@ -229,7 +320,6 @@ class RunReferenceTaskCommand(click.RichCommand):
                 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",
@@ -253,7 +343,12 @@ class RunReferenceTaskCommand(click.RichCommand):
         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()
@@ -399,7 +494,6 @@ class TaskFiles(common.FileGroup):
 
     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,
@@ -446,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

flyte/cli/_serve.py

@@ -0,0 +1,64 @@
+from typing import List
+
+import click
+
+
+@click.group("serve")
+@click.pass_context
+def serve(_: click.Context):
+    """
+    Start the specific service. For example:
+
+    ```bash
+    flyte serve connector
+    ```
+    """
+
+
+@serve.command()
+@click.option(
+    "--port",
+    default="8000",
+    is_flag=False,
+    type=int,
+    help="Grpc port for the connector service",
+)
+@click.option(
+    "--prometheus_port",
+    default="9090",
+    is_flag=False,
+    type=int,
+    help="Prometheus port for the connector service",
+)
+@click.option(
+    "--worker",
+    default="10",
+    is_flag=False,
+    type=int,
+    help="Number of workers for the grpc server",
+)
+@click.option(
+    "--timeout",
+    default=None,
+    is_flag=False,
+    type=int,
+    help="It will wait for the specified number of seconds before shutting down grpc server. It should only be used "
+    "for testing.",
+)
+@click.option(
+    "--modules",
+    required=False,
+    multiple=True,
+    type=str,
+    help="List of additional files or module that defines the connector",
+)
+@click.pass_context
+def connector(
+    _: click.Context, port: int, prometheus_port: int, worker: int, timeout: int | None, modules: List[str] | None
+):
+    """
+    Start a grpc server for the connector service.
+    """
+    from flyte.connectors import ConnectorService
+
+    ConnectorService.run(port, prometheus_port, worker, timeout, modules)

flyte/cli/_update.py

@@ -0,0 +1,37 @@
+import rich_click as click
+
+import flyte.remote as remote
+
+from . import _common as common
+
+
+@click.group(name="update")
+def update():
+    """
+    Update various flyte entities.
+    """
+
+
+@update.command("trigger", cls=common.CommandBase)
+@click.argument("name", type=str)
+@click.argument("task_name", type=str)
+@click.option("--activate/--deactivate", required=True, help="Activate or deactivate the trigger.")
+@click.pass_obj
+def trigger(cfg: common.CLIConfig, name: str, task_name: str, activate: bool, project: str | None, domain: str | None):
+    """
+    Update a trigger.
+
+    \b
+    Example usage:
+
+    ```bash
+    flyte update trigger <trigger_name> <task_name> --activate | --deactivate
+    [--project <project_name> --domain <domain_name>]
+    ```
+    """
+    cfg.init(project, domain)
+    console = common.get_console()
+    to_state = "active" if activate else "deactivate"
+    with console.status(f"Updating trigger {name} for task {task_name} to {to_state}..."):
+        remote.Trigger.update(name, task_name, activate)
+    console.print(f"Trigger updated and is set to [fuchsia]{to_state}[/fuchsia]")

flyte/cli/_user.py

@@ -0,0 +1,17 @@
+import rich_click as click
+
+import flyte.remote as remote
+
+from . import _common as common
+
+
+@click.command()
+@click.pass_obj
+def whoami(
+    cfg: common.CLIConfig,
+):
+    """Display the current user information."""
+    cfg.init()
+    console = common.get_console()
+    user_info = remote.User.get()
+    console.print(user_info.to_json())

flyte/cli/main.py

@@ -1,7 +1,7 @@
 import rich_click as click
 from typing_extensions import get_args
 
-from flyte._logging import initialize_logger, logger
+from flyte._logging import LogFormat, initialize_logger, logger
 
 from . import _common as common
 from ._abort import abort
@@ -12,7 +12,11 @@ from ._delete import delete
 from ._deploy import deploy
 from ._gen import gen
 from ._get import get
+from ._plugins import discover_and_register_plugins
 from ._run import run
+from ._serve import serve
+from ._update import update
+from ._user import whoami
 
 help_config = click.RichHelpConfiguration(
     use_markdown=True,
@@ -25,7 +29,7 @@ help_config = click.RichHelpConfiguration(
             },
             {
                 "name": "Management of various objects.",
-                "commands": ["create", "get", "delete"],
+                "commands": ["create", "get", "delete", "update"],
             },
             {
                 "name": "Build and deploy environments, tasks and images.",
@@ -35,6 +39,10 @@ help_config = click.RichHelpConfiguration(
                 "name": "Documentation generation",
                 "commands": ["gen"],
             },
+            {
+                "name": "User information",
+                "commands": ["whoami"],
+            },
         ]
     },
 )
@@ -116,6 +124,16 @@ def _verbosity_to_loglevel(verbosity: int) -> int | None:
     show_default=True,
     required=False,
 )
+@click.option(
+    "--log-format",
+    type=click.Choice(get_args(LogFormat), case_sensitive=False),
+    envvar="LOG_FORMAT",
+    default="console",
+    help="Formatting for logs, defaults to 'console' which is meant to be human readable."
+    " 'json' is meant for machine parsing.",
+    show_default=True,
+    required=False,
+)
 @click.rich_config(help_config=help_config)
 @click.pass_context
 def main(
@@ -123,6 +141,7 @@ def main(
     endpoint: str | None,
     insecure: bool,
     verbose: int,
+    log_format: LogFormat,
     org: str | None,
     config_file: str | None,
     auth_type: str | None = None,
@@ -166,8 +185,8 @@ def main(
     import flyte.config as config
 
     log_level = _verbosity_to_loglevel(verbose)
-    if log_level is not None:
-        initialize_logger(log_level)
+    if log_level is not None or log_format != "console":
+        initialize_logger(log_level=log_level, log_format=log_format)
 
     cfg = config.auto(config_file=config_file)
     if cfg.source:
@@ -175,6 +194,7 @@ def main(
 
     ctx.obj = CLIConfig(
         log_level=log_level,
+        log_format=log_format,
         endpoint=endpoint,
         insecure=insecure,
         org=org,
@@ -193,3 +213,9 @@ main.add_command(abort)  # type: ignore
 main.add_command(gen)  # type: ignore
 main.add_command(delete)  # type: ignore
 main.add_command(build)
+main.add_command(whoami)  # type: ignore
+main.add_command(update)  # type: ignore
+main.add_command(serve)  # type: ignore
+
+# Discover and register CLI plugins from installed packages
+discover_and_register_plugins(main)

flyte/config/_config.py

@@ -148,6 +148,7 @@ class ImageConfig(object):
     """
 
     builder: str | None = None
+    image_refs: typing.Dict[str, str] = field(default_factory=dict)
 
     @classmethod
     def auto(cls, config_file: typing.Optional[typing.Union[str, ConfigFile]] = None) -> "ImageConfig":
@@ -159,6 +160,7 @@ class ImageConfig(object):
         config_file = get_config_file(config_file)
         kwargs: typing.Dict[str, typing.Any] = {}
         kwargs = set_if_exists(kwargs, "builder", _internal.Image.BUILDER.read(config_file))
+        kwargs = set_if_exists(kwargs, "image_refs", _internal.Image.IMAGE_REFS.read(config_file))
         return ImageConfig(**kwargs)
 
 

flyte/config/_internal.py

@@ -70,3 +70,4 @@ class Image(object):
     """
 
     BUILDER = ConfigEntry(YamlConfigEntry("image.builder"))
+    IMAGE_REFS = ConfigEntry(YamlConfigEntry("image.image_refs"))

flyte/config/_reader.py

@@ -138,10 +138,10 @@ class ConfigFile(object):
 def _config_path_from_git_root() -> pathlib.Path | None:
     from flyte.git import config_from_root
 
-    try:
-        return config_from_root().source
-    except RuntimeError:
+    config = config_from_root()
+    if config is None:
         return None
+    return config.source
 
 
 def resolve_config_path() -> pathlib.Path | None:

flyte/connectors/__init__.py

@@ -0,0 +1,11 @@
+from ._connector import AsyncConnector, AsyncConnectorExecutorMixin, ConnectorRegistry, Resource, ResourceMeta
+from ._server import ConnectorService
+
+__all__ = [
+    "AsyncConnector",
+    "AsyncConnectorExecutorMixin",
+    "ConnectorRegistry",
+    "ConnectorService",
+    "Resource",
+    "ResourceMeta",
+]