flyte 2.0.0b25__py3-none-any.whl → 2.0.0b28__py3-none-any.whl

flyte/cli/_plugins.py

@@ -0,0 +1,208 @@
+"""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]
+    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

@@ -86,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={
@@ -119,10 +139,24 @@ class RunArguments:
             )
         },
     )
+    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]:
@@ -142,7 +176,12 @@ class RunTaskCommand(click.RichCommand):
 
     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
+            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():
@@ -153,6 +192,8 @@ class RunTaskCommand(click.RichCommand):
                 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,
             ).run.aio(self.obj, **ctx.params)
             if self.run_args.local:
                 console.print(
@@ -219,11 +260,23 @@ class TaskPerFileGroup(common.ObjectsPerFileGroup):
         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)
+        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)
+        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:
@@ -246,7 +299,12 @@ class RunReferenceTaskCommand(click.RichCommand):
 
     def invoke(self, ctx: click.Context):
         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
+            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():
@@ -284,7 +342,12 @@ class RunReferenceTaskCommand(click.RichCommand):
         import flyte.remote
         from flyte._internal.runtime.types_serde import transform_native_to_typed_interface
 
-        common.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()

flyte/cli/_serve.py

@@ -0,0 +1,154 @@
+import importlib
+import os
+import sys
+from concurrent import futures
+
+import grpc
+import rich_click as click
+from flyteidl2.service import connector_pb2
+from flyteidl2.service.connector_pb2_grpc import (
+    add_AsyncConnectorServiceServicer_to_server,
+    add_ConnectorMetadataServiceServicer_to_server,
+)
+from rich.console import Console
+from rich.table import Table
+
+
+@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, prometheus_port, worker, timeout, modules):
+    """
+    Start a grpc server for the connector service.
+    """
+    import asyncio
+
+    working_dir = os.getcwd()
+    if all(os.path.realpath(path) != working_dir for path in sys.path):
+        sys.path.append(working_dir)
+    for m in modules:
+        importlib.import_module(m)
+
+    asyncio.run(_start_grpc_server(port, prometheus_port, worker, timeout))
+
+
+async def _start_grpc_server(port: int, prometheus_port: int, worker: int, timeout: int):
+    try:
+        from flyte.connectors._server import (
+            AsyncConnectorService,
+            ConnectorMetadataService,
+        )
+    except ImportError as e:
+        raise ImportError(
+            "Flyte connector dependencies are not installed. Please install it using `pip install flyte[connector]`"
+        ) from e
+
+    click.secho("🚀 Starting the connector service...")
+    _start_http_server(prometheus_port)
+
+    print_metadata()
+
+    server = grpc.aio.server(futures.ThreadPoolExecutor(max_workers=worker))
+
+    add_AsyncConnectorServiceServicer_to_server(AsyncConnectorService(), server)
+    add_ConnectorMetadataServiceServicer_to_server(ConnectorMetadataService(), server)
+    _start_health_check_server(server, worker)
+
+    server.add_insecure_port(f"[::]:{port}")
+    await server.start()
+    await server.wait_for_termination(timeout)
+
+
+def _start_http_server(prometheus_port: int):
+    try:
+        from prometheus_client import start_http_server
+
+        click.secho("Starting up the server to expose the prometheus metrics...")
+        start_http_server(prometheus_port)
+    except ImportError as e:
+        click.secho(f"Failed to start the prometheus server with error {e}", fg="red")
+
+
+def _start_health_check_server(server: grpc.Server, worker: int):
+    try:
+        from grpc_health.v1 import health, health_pb2, health_pb2_grpc
+
+        health_servicer = health.HealthServicer(
+            experimental_non_blocking=True,
+            experimental_thread_pool=futures.ThreadPoolExecutor(max_workers=worker),
+        )
+
+        for service in connector_pb2.DESCRIPTOR.services_by_name.values():
+            health_servicer.set(service.full_name, health_pb2.HealthCheckResponse.SERVING)
+        health_servicer.set(health.SERVICE_NAME, health_pb2.HealthCheckResponse.SERVING)
+
+        health_pb2_grpc.add_HealthServicer_to_server(health_servicer, server)
+
+    except ImportError as e:
+        click.secho(f"Failed to start the health check servicer with error {e}", fg="red")
+
+
+def print_metadata():
+    from flyte.connectors import ConnectorRegistry
+
+    connectors = ConnectorRegistry.list_connectors()
+
+    table = Table(title="Connector Metadata")
+    table.add_column("Connector Name", style="cyan", no_wrap=True)
+    table.add_column("Support Task Types", style="cyan")
+
+    for connector in connectors:
+        categories = ""
+        for category in connector.supported_task_categories:
+            categories += f"{category.name} ({category.version}) "
+        table.add_row(connector.name, categories)
+
+    console = Console()
+    console.print(table)

flyte/cli/main.py

@@ -12,7 +12,9 @@ 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
 
@@ -201,3 +203,7 @@ 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/connectors/__init__.py

@@ -0,0 +1,3 @@
+from ._connector import AsyncConnector, ConnectorRegistry, Resource, ResourceMeta
+
+__all__ = ["AsyncConnector", "ConnectorRegistry", "Resource", "ResourceMeta"]