_prefetch.py

"""
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(
    "--allow-pattern",
    "allow_patterns",
    type=str,
    multiple=True,
    help=(
        "Glob pattern selecting which repo files to prefetch, e.g. `*Q4_K_M*` to pull one GGUF "
        "quant out of a repo that ships many. Can be specified multiple times. Omit to prefetch "
        "the whole repo. Ignored when `--shard-config` is set."
    ),
)
@click.option(
    "--ignore-pattern",
    "ignore_patterns",
    type=str,
    multiple=True,
    help="Glob pattern excluded from the prefetch, applied after `--allow-pattern`. Can be specified multiple times.",
)
@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,
    allow_patterns: tuple[str, ...],
    ignore_patterns: tuple[str, ...],
    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 \\
        --gpu 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,
        allow_patterns=list(allow_patterns) or None,
        ignore_patterns=list(ignore_patterns) or None,
        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}")