main.py

import rich_click as click
from typing_extensions import get_args

import flyte
from flyte._logging import _LOG_LEVEL_MAP, LogFormat, initialize_logger, logger

from . import _common as common
from ._abort import abort
from ._build import build
from ._common import CLIConfig
from ._create import create
from ._delete import delete
from ._deploy import deploy
from ._edit import edit
from ._gen import gen
from ._get import get
from ._plugins import discover_and_register_plugins
from ._prefetch import prefetch
from ._proxy import proxy
from ._rerun import rerun
from ._run import run
from ._serve import serve
from ._signal import signal
from ._start import start
from ._stop import stop
from ._update import update
from ._user import whoami

help_config = click.RichHelpConfiguration(
    use_markdown=True,
    use_markdown_emoji=True,
    command_groups={
        "flyte": [
            {
                "name": "Run and stop tasks",
                "commands": ["run", "rerun", "abort", "signal"],
            },
            {
                "name": "Serve Apps",
                "commands": ["serve"],
            },
            {
                "name": "Management of various objects.",
                "commands": ["create", "get", "delete", "update"],
            },
            {
                "name": "Settings management.",
                "commands": ["edit"],
            },
            {
                "name": "Build and deploy environments, tasks and images.",
                "commands": ["build", "deploy"],
            },
            {
                "name": "Prefetch artifacts from remote registries.",
                "commands": ["prefetch"],
            },
            {
                "name": "Documentation generation",
                "commands": ["gen"],
            },
            {
                "name": "User information",
                "commands": ["whoami"],
            },
        ]
    },
)


def _verbosity_to_loglevel(verbosity: int) -> int | None:
    """
    Converts a verbosity level from the CLI to a logging level.

    Args:
        verbosity: verbosity level from the CLI

    Returns:
        logging level
    """
    import logging

    match verbosity:
        case 0:
            return None
        case 1:
            return logging.WARNING
        case 2:
            return logging.INFO
        case _:
            return logging.DEBUG


@click.group(cls=click.RichGroup)
@click.version_option(
    message=f"Flyte SDK version: {flyte.version()}",
)
@click.option(
    "--endpoint",
    type=str,
    required=False,
    help="The endpoint to connect to. This will override any configuration file and simply use `pkce` to connect.",
)
@click.option(
    "--insecure",
    is_flag=True,
    required=False,
    help="Use an insecure connection to the endpoint. If not specified, the CLI will use TLS.",
    type=bool,
    default=None,
    show_default=True,
)
@click.option(
    "--image-builder",
    "--builder",
    type=click.Choice(["local", "remote"]),
    default=None,
    help="Image builder to use for building images. Overrides the config file setting."
    " If not specified, the builder from the config file (image.builder) is used,"
    " falling back to `local`.",
    show_default=True,
    required=False,
)
@click.option(
    "--auth-type",
    type=click.Choice(common.ALL_AUTH_OPTIONS, case_sensitive=False),
    default=None,
    help="Authentication type to use for the Flyte backend. Defaults to `pkce`.",
    show_default=True,
    required=False,
)
@click.option(
    "-v",
    "--verbose",
    required=False,
    help="Show verbose messages and exception traces. Repeating multiple times increases the verbosity (e.g., -vvv).",
    count=True,
    default=0,
    type=int,
)
@click.option(
    "--org",
    type=str,
    required=False,
    help="The organization to which the command applies.",
)
@click.option(
    "-c",
    "--config",
    "config_file",
    required=False,
    type=click.Path(exists=True, dir_okay=False),
    help="Path to the configuration file to use. If not specified, the default configuration file is used.",
)
@click.option(
    "--output-format",
    "-of",
    type=click.Choice(get_args(common.OutputFormat), case_sensitive=False),
    default="table",
    help="Output format for commands that support it. Defaults to `table`.",
    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.option(
    "--user-log-level",
    type=click.Choice(["debug", "info", "warning", "error", "critical"], case_sensitive=False),
    envvar="USER_LOG_LEVEL",
    default="info",
    show_default=True,
    help="Log level for user task logs. Independent of the internal Flyte log level (-v).",
    required=False,
)
@click.option(
    "--reset-root-logger",
    is_flag=True,
    required=False,
    help="If set, the root logger will be reset to use Flyte logging style",
    type=bool,
    default=False,
    show_default=True,
)
@click.option(
    "--no-progress",
    is_flag=True,
    required=False,
    help="Disable the animated progress spinner — useful in CI / non-interactive logs.",
    type=bool,
    default=False,
    show_default=True,
)
@click.rich_config(help_config=help_config)
@click.pass_context
def main(
    ctx: click.Context,
    endpoint: str | None,
    insecure: bool,
    image_builder: str | None,
    verbose: int,
    log_format: LogFormat,
    reset_root_logger: bool,
    org: str | None,
    config_file: str | None,
    auth_type: str | None = None,
    output_format: common.OutputFormat = "table",
    user_log_level: str = "info",
    no_progress: bool = False,
):
    """
    The Flyte CLI is the command line interface for working with the Flyte SDK and backend.

    It follows a simple verb/noun structure,
    where the top-level commands are verbs that describe the action to be taken,
    and the subcommands are nouns that describe the object of the action.

    The root command can be used to configure the CLI for persistent settings,
    such as the endpoint, organization, and verbosity level.

    Set endpoint and organization:

    ```bash
    $ flyte --endpoint <endpoint> --org <org> get project <project_name>
    ```

    Increase verbosity level (This is useful for debugging,
    this will show more logs and exception traces):

    ```bash
    $ flyte -vvv get logs <run-name>
    ```

    Override the default config file:

    ```bash
    $ flyte --config /path/to/config.yaml run ...
    ```

    * [Documentation](https://www.union.ai/docs/flyte/user-guide/)
    * [GitHub](https://github.com/flyteorg/flyte): Please leave a star if you like Flyte!
    * [Slack](https://slack.flyte.org): Join the community and ask questions.
    * [Issues](https://github.com/flyteorg/flyte/issues)

    """
    import flyte.config as config

    log_level = _verbosity_to_loglevel(verbose)
    user_log_level_int = _LOG_LEVEL_MAP[user_log_level.lower()]
    initialize_logger(
        log_level=log_level,
        log_format=log_format,
        enable_rich=True,
        reset_root_logger=reset_root_logger,
        user_log_level=user_log_level_int,
    )

    cfg = config.auto(config_file=config_file)
    if cfg.source:
        logger.debug(f"Using config file discovered at location `{cfg.source.absolute()}`")

    ctx.obj = CLIConfig(
        log_level=log_level,
        log_format=log_format,
        reset_root_logger=reset_root_logger,
        user_log_level=user_log_level_int,
        endpoint=endpoint,
        insecure=insecure,
        image_builder=image_builder,
        org=org,
        config=cfg,
        ctx=ctx,
        auth_type=auth_type,
        output_format=output_format,
        no_progress=no_progress,
    )

    from flyte._status import set_output_mode

    set_output_mode("rich" if output_format == "table" else "plain")


main.add_command(run)
main.add_command(rerun)
main.add_command(deploy)
main.add_command(get)  # type: ignore
main.add_command(create)  # type: ignore
main.add_command(abort)  # type: ignore
main.add_command(signal)  # 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
main.add_command(start)  # type: ignore
main.add_command(stop)  # type: ignore
main.add_command(prefetch)  # type: ignore
main.add_command(edit)  # type: ignore
main.add_command(proxy)  # type: ignore

# Discover and register CLI plugins from installed packages
discover_and_register_plugins(main)