flyte 0.2.0b1__py3-none-any.whl → 2.0.0b46__py3-none-any.whl

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/_prefetch.py

@@ -0,0 +1,292 @@
+"""
+CLI commands for prefetching artifacts from remote registries.
+"""
+
+import typing
+from pathlib import Path
+
+import rich_click as click
+from rich.console import Console
+
+from flyte._resources import Accelerators
+from flyte.cli._common import CommandBase
+
+# Get all valid accelerator choices from the Accelerators literal type
+ACCELERATOR_CHOICES = list(typing.get_args(Accelerators))
+
+
+@click.group(name="prefetch")
+def prefetch():
+    """
+    Prefetch artifacts from remote registries.
+
+    These commands help you download and prefetch artifacts like HuggingFace models
+    to your Flyte storage for faster access during task execution.
+    """
+
+
+@prefetch.command(name="hf-model", cls=CommandBase)
+@click.argument("repo", type=str)
+@click.option(
+    "--raw-data-path",
+    type=str,
+    required=False,
+    default=None,
+    help=(
+        "Object store path to store the model. If not provided, the model will be stored using the default path "
+        "generated by Flyte storage layer."
+    ),
+)
+@click.option(
+    "--artifact-name",
+    type=str,
+    required=False,
+    default=None,
+    help=(
+        "Artifact name to use for the stored model. Must only contain alphanumeric characters, "
+        "underscores, and hyphens. If not provided, the repo name will be used (replacing '.' with '-')."
+    ),
+)
+@click.option(
+    "--architecture",
+    type=str,
+    help="Model architecture, as given in HuggingFace config.json.",
+)
+@click.option(
+    "--task",
+    default="auto",
+    type=str,
+    help=(
+        "Model task, e.g., 'generate', 'classify', 'embed', 'score', etc. "
+        "Refer to vLLM docs. 'auto' will try to discover this automatically."
+    ),
+)
+@click.option(
+    "--modality",
+    type=str,
+    multiple=True,
+    default=("text",),
+    help="Modalities supported by the model, e.g., 'text', 'image', 'audio', 'video'. Can be specified multiple times.",
+)
+@click.option(
+    "--format",
+    "serial_format",
+    type=str,
+    help="Model serialization format, e.g., safetensors, onnx, torchscript, joblib, etc.",
+)
+@click.option(
+    "--model-type",
+    type=str,
+    help=(
+        "Model type, e.g., 'transformer', 'xgboost', 'custom', etc. "
+        "For HuggingFace models, this is auto-determined from config.json['model_type']."
+    ),
+)
+@click.option(
+    "--short-description",
+    type=str,
+    help="Short description of the model.",
+)
+@click.option(
+    "--force",
+    type=int,
+    default=0,
+    help="Force store of the model. Increment value (--force=1, --force=2, ...) to force a new store.",
+)
+@click.option(
+    "--wait",
+    is_flag=True,
+    help="Wait for the model to be stored before returning.",
+)
+@click.option(
+    "--hf-token-key",
+    type=str,
+    default="HF_TOKEN",
+    help=(
+        "Name of the Flyte secret containing your HuggingFace token. "
+        "Note: This is not the HuggingFace token itself, but the name of the "
+        "secret in the Flyte secret store."
+    ),
+    show_default=True,
+)
+@click.option(
+    "--cpu",
+    type=str,
+    default="2",
+    help="CPU request for the prefetch task (e.g., '2', '4', '2,4' for 2-4 CPUs).",
+)
+@click.option(
+    "--mem",
+    type=str,
+    default="8Gi",
+    help="Memory request for the prefetch task (e.g., '16Gi', '64Gi', '16Gi,64Gi' for 16-64GB).",
+)
+@click.option(
+    "--gpu",
+    type=click.Choice(ACCELERATOR_CHOICES),
+    default=None,
+    help=(
+        "The gpu to use for downloading and (optionally) sharding the model. "
+        "Format: '{type}:{quantity}' (e.g., 'A100:8', 'L4:1')."
+    ),
+)
+@click.option(
+    "--disk",
+    type=str,
+    default="50Gi",
+    help="Disk storage request for the prefetch task (e.g., '100Gi', '500Gi').",
+)
+@click.option(
+    "--shm",
+    type=str,
+    default=None,
+    help="Shared memory request for the prefetch task (e.g., '100Gi', 'auto').",
+)
+@click.option(
+    "--shard-config",
+    type=click.Path(exists=True, path_type=Path),
+    help=(
+        "Path to a YAML file containing sharding configuration. "
+        "The file should have 'engine' (currently only 'vllm') and 'args' keys."
+    ),
+)
+@click.pass_obj
+def hf_model(
+    cfg,
+    repo: str,
+    raw_data_path: str | None,
+    artifact_name: str | None,
+    architecture: str | None,
+    task: str,
+    modality: tuple[str, ...],
+    serial_format: str | None,
+    model_type: str | None,
+    short_description: str | None,
+    force: int,
+    wait: bool,
+    hf_token_key: str,
+    cpu: str | None,
+    mem: str | None,
+    disk: str | None,
+    gpu: Accelerators | None,
+    shm: str | None,
+    shard_config: Path | None,
+    project: str | None,
+    domain: str | None,
+):
+    """
+    Prefetch a HuggingFace model to Flyte storage.
+
+    Downloads a model from the HuggingFace Hub and prefetches it to your configured
+    Flyte storage backend. This is useful for:
+
+    - Pre-fetching large models before running inference tasks
+    - Sharding models for tensor-parallel inference
+    - Avoiding repeated downloads during development
+
+    **Basic Usage:**
+
+    ```bash
+    $ flyte prefetch hf-model meta-llama/Llama-2-7b-hf --hf-token-key HF_TOKEN
+    ```
+
+    **With Sharding:**
+
+    Create a shard config file (shard_config.yaml):
+
+    ```yaml
+    engine: vllm
+    args:
+      tensor_parallel_size: 8
+      dtype: auto
+      trust_remote_code: true
+    ```
+
+    Then run:
+
+    ```bash
+    $ flyte prefetch hf-model meta-llama/Llama-2-70b-hf \\
+        --shard-config shard_config.yaml \\
+        --accelerator A100:8 \\
+        --hf-token-key HF_TOKEN
+    ```
+
+    **Wait for Completion:**
+
+    ```bash
+    $ flyte prefetch hf-model meta-llama/Llama-2-7b-hf --wait
+    ```
+    """
+    import yaml
+
+    from flyte._resources import Resources
+    from flyte.cli._run import initialize_config
+    from flyte.prefetch import ShardConfig, VLLMShardArgs
+    from flyte.prefetch import hf_model as prefetch_hf_model
+
+    # Initialize flyte config
+    cfg = initialize_config(
+        cfg.ctx,
+        project or cfg.config.task.project,
+        domain or cfg.config.task.domain,
+    )
+
+    # Parse shard config if provided
+    parsed_shard_config = None
+    if shard_config is not None:
+        with shard_config.open() as f:
+            shard_config_dict = yaml.safe_load(f)
+            args_dict = shard_config_dict.get("args", {})
+            parsed_shard_config = ShardConfig(
+                engine=shard_config_dict.get("engine", "vllm"),
+                args=VLLMShardArgs(**args_dict),
+            )
+
+    console = Console()
+
+    console.print("[bold green]Starting model prefetch task...")
+
+    # Parse cpu and mem for range syntax (e.g., "2, 4" -> ("2", "4"))
+    parsed_cpu: str | tuple[str, str] | None = cpu
+    if cpu is not None:
+        cpu_parts = cpu.split(", ")
+        if len(cpu_parts) > 1:
+            parsed_cpu = (cpu_parts[0], cpu_parts[1])
+
+    parsed_mem: str | tuple[str, str] | None = mem
+    if mem is not None:
+        mem_parts = mem.split(", ")
+        if len(mem_parts) > 1:
+            parsed_mem = (mem_parts[0], mem_parts[1])
+
+    run = prefetch_hf_model(
+        repo=repo,
+        raw_data_path=raw_data_path,
+        artifact_name=artifact_name,
+        architecture=architecture,
+        task=task,
+        modality=modality,
+        serial_format=serial_format,
+        model_type=model_type,
+        short_description=short_description,
+        shard_config=parsed_shard_config,
+        hf_token_key=hf_token_key,
+        resources=Resources(cpu=parsed_cpu, memory=parsed_mem, disk=disk, gpu=gpu, shm=shm),
+        force=force,
+    )
+
+    url = run.url
+    console.print(
+        f"🔄 Started run {run.name} to prefetch model from HuggingFace repo [bold]{repo}[/bold].\n"
+        f"   Check the console for status at [link={url}]{url}[/link]"
+    )
+
+    if wait:
+        run.wait()
+        try:
+            model_path = run.outputs()[0].path
+            console.print("\n✅ Model prefetched successfully!")
+            console.print(f"Remote path: [cyan]{model_path}[/cyan]")
+        except Exception as e:
+            console.print("\n❌ Model prefetch failed!")
+            console.print(f"Error: {e}")