flyte 0.1.0__py3-none-any.whl → 0.2.0a0__py3-none-any.whl

flyte/cli/_run.py

@@ -0,0 +1,231 @@
+from __future__ import annotations
+
+import asyncio
+import inspect
+from dataclasses import dataclass, field, fields
+from pathlib import Path
+from types import ModuleType
+from typing import Any, Dict, List, cast
+
+import click
+from click import Context, Parameter
+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 ._params import to_click_option
+
+
+@dataclass
+class RunArguments:
+    project: str = field(
+        default=cast(str, common.PROJECT_OPTION.default), metadata={"click.option": common.PROJECT_OPTION}
+    )
+    domain: str = field(
+        default=cast(str, common.DOMAIN_OPTION.default), metadata={"click.option": common.DOMAIN_OPTION}
+    )
+    local: bool = field(
+        default=False,
+        metadata={
+            "click.option": click.Option(
+                ["--local"],
+                is_flag=True,
+                help="Run the task locally",
+            )
+        },
+    )
+    copy_style: CopyFiles = field(
+        default="loaded_modules",
+        metadata={
+            "click.option": click.Option(
+                ["--copy-style"],
+                type=click.Choice(get_args(CopyFiles)),
+                default="loaded_modules",
+                help="Copy style to use when running the task",
+            )
+        },
+    )
+    name: str | None = field(
+        default=None,
+        metadata={
+            "click.option": click.Option(
+                ["--name"],
+                type=str,
+                help="Name of the run. If not provided, a random name will be generated.",
+            )
+        },
+    )
+    follow: bool = field(
+        default=True,
+        metadata={
+            "click.option": click.Option(
+                ["--follow", "-f"],
+                is_flag=True,
+                default=False,
+                help="Wait and watch logs for the parent action. If not provided, the CLI will exit after "
+                "successfully launching a remote execution with a link to the UI.",
+            )
+        },
+    )
+
+    @classmethod
+    def from_dict(cls, d: Dict[str, Any]) -> RunArguments:
+        return cls(**d)
+
+    @classmethod
+    def options(cls) -> List[click.Option]:
+        """
+        Return the set of base parameters added to run subcommand.
+        """
+        return [common.get_option_from_metadata(f.metadata) for f in fields(cls) if f.metadata]
+
+
+class RunTaskCommand(click.Command):
+    def __init__(self, obj_name: str, obj: Any, run_args: RunArguments, *args, **kwargs):
+        self.obj_name = obj_name
+        self.obj = cast(TaskTemplate, obj)
+        self.run_args = run_args
+        kwargs.pop("name", None)
+        super().__init__(obj_name, *args, **kwargs)
+
+    def invoke(self, ctx: Context):
+        obj: CLIConfig = ctx.obj
+        if obj is None:
+            import flyte.config
+
+            obj = CLIConfig(flyte.config.auto(), ctx)
+
+        obj.init(self.run_args.project, self.run_args.domain)
+
+        async def _run():
+            import flyte
+
+            r = flyte.with_runcontext(
+                copy_style=self.run_args.copy_style,
+                mode="local" if self.run_args.local else "remote",
+                name=self.run_args.name,
+            ).run(self.obj, **ctx.params)
+            if isinstance(r, Run) and r.action is not None:
+                console = Console()
+                console.print(
+                    common.get_panel(
+                        "Run",
+                        f"[green bold]Created Run: {r.name} [/green bold] "
+                        f"(Project: {r.action.action_id.run.project}, Domain: {r.action.action_id.run.domain})\n"
+                        f"➡️  [blue bold]{r.url}[/blue bold]",
+                    )
+                )
+                if self.run_args.follow:
+                    console.print(
+                        "[dim]Log streaming enabled, will wait for task to start running "
+                        "and log stream to be available[/dim]"
+                    )
+                    await r.show_logs(max_lines=30, show_ts=True, raw=False)
+
+        asyncio.run(_run())
+
+    def get_params(self, ctx: Context) -> List[Parameter]:
+        # Note this function may be called multiple times by click.
+        task = self.obj
+        from .._internal.runtime.types_serde import transform_native_to_typed_interface
+
+        interface = transform_native_to_typed_interface(task.native_interface)
+        if interface is None:
+            return super().get_params(ctx)
+        inputs_interface = task.native_interface.inputs
+
+        params: List[Parameter] = []
+        for name, var in interface.inputs.variables.items():
+            default_val = None
+            if inputs_interface[name][1] is not inspect._empty:
+                default_val = inputs_interface[name][1]
+            params.append(to_click_option(name, var, inputs_interface[name][0], default_val))
+
+        self.params = params
+        return super().get_params(ctx)
+
+
+class TaskPerFileGroup(common.ObjectsPerFileGroup):
+    """
+    Group that creates a command for each task in the current directory that is not __init__.py.
+    """
+
+    def __init__(self, filename: Path, run_args: RunArguments, *args, **kwargs):
+        args = (filename, *args)
+        super().__init__(*args, **kwargs)
+        self.run_args = run_args
+
+    def _filter_objects(self, module: ModuleType) -> Dict[str, Any]:
+        return {k: v for k, v in module.__dict__.items() if isinstance(v, TaskTemplate)}
+
+    def _get_command_for_obj(self, ctx: click.Context, obj_name: str, obj: Any) -> click.Command:
+        obj = cast(TaskTemplate, obj)
+        return RunTaskCommand(
+            obj_name=obj_name,
+            obj=obj,
+            help=obj.docs.__help__str__() if obj.docs else None,
+            run_args=self.run_args,
+        )
+
+
+class TaskFiles(common.FileGroup):
+    """
+    Group that creates a command for each file in the current directory that is not __init__.py.
+    """
+
+    common_options_enabled = False
+
+    def __init__(
+        self,
+        *args,
+        directory: Path | None = None,
+        **kwargs,
+    ):
+        if "params" not in kwargs:
+            kwargs["params"] = []
+        kwargs["params"].extend(RunArguments.options())
+        super().__init__(*args, directory=directory, **kwargs)
+
+    def get_command(self, ctx, filename):
+        run_args = RunArguments.from_dict(ctx.params)
+        fp = Path(filename)
+        if not fp.exists():
+            raise click.BadParameter(f"File {filename} does not exist")
+        if fp.is_dir():
+            return TaskFiles(directory=fp)
+        return TaskPerFileGroup(
+            filename=Path(filename),
+            run_args=run_args,
+            name=filename,
+            help=f"Run, functions decorated with `env.task` in {filename}",
+        )
+
+
+run = TaskFiles(
+    name="run",
+    help="""
+Run a task from a python file.
+
+Example usage:
+```bash
+flyte run --name examples/basics/hello.py my_task --arg1 value1 --arg2 value2
+```
+Note: all arguments for the run command are provided right after the `run` command and before the file name.
+
+You can also specify the project and domain using the `--project` and `--domain` options, respectively. These
+options can be set in the config file or passed as command line arguments.
+
+Note: The arguments for the task are provided after the task name and can be retrieved using `--help`
+Example:
+```bash
+flyte run --name examples/basics/hello.py my_task --help
+```
+
+To run a task locally, use the `--local` flag. This will run the task in the local environment instead of the remote
+ Flyte environment.
+""",
+)

flyte/cli/main.py

@@ -0,0 +1,166 @@
+import rich_click as click
+
+from flyte._logging import initialize_logger, logger
+
+from ._abort import abort
+from ._common import CLIConfig
+from ._create import create
+from ._delete import delete
+from ._deploy import deploy
+from ._gen import gen
+from ._get import get
+from ._run import run
+
+help_config = click.RichHelpConfiguration(
+    use_markdown=True,
+    use_markdown_emoji=True,
+    command_groups={
+        "flyte": [
+            {
+                "name": "Run and stop tasks",
+                "commands": ["run", "abort"],
+            },
+            {
+                "name": "Management",
+                "commands": ["create", "deploy", "get", "delete"],
+            },
+            {
+                "name": "Documentation generation",
+                "commands": ["gen"],
+            },
+        ]
+    },
+)
+
+
+def _verbosity_to_loglevel(verbosity: int) -> int | None:
+    """
+    Converts a verbosity level from the CLI to a logging level.
+
+    :param verbosity: verbosity level from the CLI
+    :return: 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.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(
+    "-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),
+    help="Path to the configuration file to use. If not specified, the default configuration file is used.",
+)
+@click.rich_config(help_config=help_config)
+@click.pass_context
+def main(
+    ctx: click.Context,
+    endpoint: str | None,
+    insecure: bool,
+    verbose: int,
+    org: str | None,
+    config_file: str | None,
+):
+    """
+    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)
+    if log_level is not None:
+        initialize_logger(log_level)
+
+    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,
+        endpoint=endpoint,
+        insecure=insecure,
+        org=org,
+        config=cfg,
+        ctx=ctx,
+    )
+
+
+main.add_command(run)
+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(gen)  # type: ignore
+main.add_command(delete)  # type: ignore

flyte/config/__init__.py

@@ -0,0 +1,3 @@
+from flyte.config._config import Config, auto, get_config_file, set_if_exists
+
+__all__ = ["Config", "auto", "get_config_file", "set_if_exists"]

flyte/config/_config.py

@@ -0,0 +1,216 @@
+from __future__ import annotations
+
+import dataclasses
+import os
+import pathlib
+import typing
+from dataclasses import dataclass, field
+from typing import TYPE_CHECKING
+
+import rich.repr
+
+from flyte._logging import logger
+from flyte.config import _internal
+from flyte.config._reader import ConfigFile, get_config_file, read_file_if_exists
+
+_all__ = ["ConfigFile", "PlatformConfig", "TaskConfig"]
+
+if TYPE_CHECKING:
+    from flyte.remote._client.auth import AuthType
+
+
+@rich.repr.auto
+@dataclass(init=True, repr=True, eq=True, frozen=True)
+class PlatformConfig(object):
+    """
+    This object contains the settings to talk to a Flyte backend (the DNS location of your Admin server basically).
+
+    :param endpoint: DNS for Flyte backend
+    :param insecure: Whether or not to use SSL
+    :param insecure_skip_verify: Whether to skip SSL certificate verification
+    :param console_endpoint: endpoint for console if different from Flyte backend
+    :param command: This command is executed to return a token using an external process
+    :param proxy_command: This command is executed to return a token for proxy authorization using an external process
+    :param client_id: This is the public identifier for the app which handles authorization for a Flyte deployment.
+      More details here: https://www.oauth.com/oauth2-servers/client-registration/client-id-secret/.
+    :param client_credentials_secret: Used for service auth, which is automatically called during pyflyte. This will
+      allow the Flyte engine to read the password directly from the environment variable. Note that this is
+      less secure! Please only use this if mounting the secret as a file is impossible
+    :param scopes: List of scopes to request. This is only applicable to the client credentials flow
+    :param auth_mode: The OAuth mode to use. Defaults to pkce flow
+    :param ca_cert_file_path: [optional] str Root Cert to be loaded and used to verify admin
+    :param http_proxy_url: [optional] HTTP Proxy to be used for OAuth requests
+    """
+
+    endpoint: str | None = None
+    insecure: bool = False
+    insecure_skip_verify: bool = False
+    ca_cert_file_path: typing.Optional[str] = None
+    console_endpoint: typing.Optional[str] = None
+    command: typing.Optional[typing.List[str]] = None
+    proxy_command: typing.Optional[typing.List[str]] = None
+    client_id: typing.Optional[str] = None
+    client_credentials_secret: typing.Optional[str] = None
+    scopes: typing.List[str] = field(default_factory=list)
+    auth_mode: "AuthType" = "Pkce"
+    audience: typing.Optional[str] = None
+    rpc_retries: int = 3
+    http_proxy_url: typing.Optional[str] = None
+
+    @classmethod
+    def auto(cls, config_file: typing.Optional[typing.Union[str, ConfigFile]] = None) -> "PlatformConfig":
+        """
+        Reads from a config file, and overrides from Environment variables. Refer to ConfigEntry for details
+        :param config_file:
+        :return:
+        """
+
+        config_file = get_config_file(config_file)
+        kwargs: typing.Dict[str, typing.Any] = {}
+        kwargs = set_if_exists(kwargs, "insecure", _internal.Platform.INSECURE.read(config_file))
+        kwargs = set_if_exists(
+            kwargs, "insecure_skip_verify", _internal.Platform.INSECURE_SKIP_VERIFY.read(config_file)
+        )
+        kwargs = set_if_exists(kwargs, "ca_cert_file_path", _internal.Platform.CA_CERT_FILE_PATH.read(config_file))
+        kwargs = set_if_exists(kwargs, "command", _internal.Credentials.COMMAND.read(config_file))
+        kwargs = set_if_exists(kwargs, "proxy_command", _internal.Credentials.PROXY_COMMAND.read(config_file))
+        kwargs = set_if_exists(kwargs, "client_id", _internal.Credentials.CLIENT_ID.read(config_file))
+
+        is_client_secret = False
+        client_credentials_secret = read_file_if_exists(
+            _internal.Credentials.CLIENT_CREDENTIALS_SECRET_LOCATION.read(config_file)
+        )
+        if client_credentials_secret:
+            is_client_secret = True
+            if client_credentials_secret.endswith("\n"):
+                logger.info("Newline stripped from client secret")
+                client_credentials_secret = client_credentials_secret.strip()
+        kwargs = set_if_exists(
+            kwargs,
+            "client_credentials_secret",
+            client_credentials_secret,
+        )
+
+        client_credentials_secret_env_var = _internal.Credentials.CLIENT_CREDENTIALS_SECRET_ENV_VAR.read(config_file)
+        if client_credentials_secret_env_var:
+            client_credentials_secret = os.getenv(client_credentials_secret_env_var)
+            if client_credentials_secret:
+                is_client_secret = True
+        kwargs = set_if_exists(kwargs, "client_credentials_secret", client_credentials_secret)
+        kwargs = set_if_exists(kwargs, "scopes", _internal.Credentials.SCOPES.read(config_file))
+        kwargs = set_if_exists(kwargs, "auth_mode", _internal.Credentials.AUTH_MODE.read(config_file))
+        if is_client_secret:
+            kwargs = set_if_exists(kwargs, "auth_mode", "ClientSecret")
+        kwargs = set_if_exists(kwargs, "endpoint", _internal.Platform.URL.read(config_file))
+        kwargs = set_if_exists(kwargs, "console_endpoint", _internal.Platform.CONSOLE_ENDPOINT.read(config_file))
+
+        kwargs = set_if_exists(kwargs, "http_proxy_url", _internal.Platform.HTTP_PROXY_URL.read(config_file))
+        return PlatformConfig(**kwargs)
+
+    def replace(self, **kwargs: typing.Any) -> "PlatformConfig":
+        """
+        Returns a new PlatformConfig instance with the values from the kwargs overriding the current instance.
+        """
+        return dataclasses.replace(self, **kwargs)
+
+    @classmethod
+    def for_endpoint(cls, endpoint: str, insecure: bool = False) -> "PlatformConfig":
+        return PlatformConfig(endpoint=endpoint, insecure=insecure)
+
+
+@rich.repr.auto
+@dataclass(init=True, repr=True, eq=True, frozen=True)
+class TaskConfig(object):
+    org: str | None = None
+    project: str | None = None
+    domain: str | None = None
+
+    @classmethod
+    def auto(cls, config_file: typing.Optional[typing.Union[str, ConfigFile]] = None) -> "TaskConfig":
+        """
+        Reads from a config file, and overrides from Environment variables. Refer to ConfigEntry for details
+        :param config_file:
+        :return:
+        """
+        config_file = get_config_file(config_file)
+        kwargs: typing.Dict[str, typing.Any] = {}
+        kwargs = set_if_exists(kwargs, "org", _internal.Task.ORG.read(config_file))
+        kwargs = set_if_exists(kwargs, "project", _internal.Task.PROJECT.read(config_file))
+        kwargs = set_if_exists(kwargs, "domain", _internal.Task.DOMAIN.read(config_file))
+        return TaskConfig(**kwargs)
+
+
+@rich.repr.auto
+@dataclass(init=True, repr=True, eq=True, frozen=True)
+class Config(object):
+    """
+    This the parent configuration object and holds all the underlying configuration object types. An instance of
+    this object holds all the config necessary to
+
+    1. Interactive session with Flyte backend
+    2. Some parts are required for Serialization, for example Platform Config is not required
+    3. Runtime of a task
+    """
+
+    platform: PlatformConfig = field(default=PlatformConfig())
+    task: TaskConfig = field(default=TaskConfig())
+    source: pathlib.Path | None = None
+
+    def with_params(
+        self,
+        platform: PlatformConfig | None = None,
+        task: TaskConfig | None = None,
+    ) -> "Config":
+        return Config(
+            platform=platform or self.platform,
+            task=task or self.task,
+        )
+
+    @classmethod
+    def auto(cls, config_file: typing.Union[str, ConfigFile, None] = None) -> "Config":
+        """
+        Automatically constructs the Config Object. The order of precedence is as follows
+          1. first try to find any env vars that match the config vars specified in the FLYTE_CONFIG format.
+          2. If not found in environment then values ar read from the config file
+          3. If not found in the file, then the default values are used.
+
+        :param config_file: file path to read the config from, if not specified default locations are searched
+        :return: Config
+        """
+        config_file = get_config_file(config_file)
+        if config_file is None:
+            logger.debug("No config file found, using default values")
+            return Config()
+        return Config(
+            platform=PlatformConfig.auto(config_file), task=TaskConfig.auto(config_file), source=config_file.path
+        )
+
+
+def set_if_exists(d: dict, k: str, val: typing.Any) -> dict:
+    """
+    Given a dict ``d`` sets the key ``k`` with value of config ``v``, if the config value ``v`` is set
+    and return the updated dictionary.
+    """
+    exists = isinstance(val, bool) or bool(val is not None and val)
+    if exists:
+        d[k] = val
+    return d
+
+
+def auto(config_file: typing.Union[str, ConfigFile, None] = None) -> Config:
+    """
+    Automatically constructs the Config Object. The order of precedence is as follows
+      1. If specified, read the config from the provided file path.
+      2. If not specified, the config file is searched in the default locations.
+            a. ./config.yaml if it exists  (current working directory)
+            b. `UCTL_CONFIG` environment variable
+            c. `FLYTECTL_CONFIG` environment variable
+            d. ~/.union/config.yaml if it exists
+            e. ~/.flyte/config.yaml if it exists
+    3. If any value is not found in the config file, the default value is used.
+    4. For any value there are environment variables that match the config variable names, those will override
+
+    :param config_file: file path to read the config from, if not specified default locations are searched
+    :return: Config
+    """
+    return Config.auto(config_file)

flyte/config/_internal.py

@@ -0,0 +1,64 @@
+from flyte.config._reader import ConfigEntry, YamlConfigEntry
+
+
+class Platform(object):
+    URL = ConfigEntry(YamlConfigEntry("admin.endpoint"))
+    INSECURE = ConfigEntry(YamlConfigEntry("admin.insecure", bool))
+    INSECURE_SKIP_VERIFY = ConfigEntry(YamlConfigEntry("admin.insecureSkipVerify", bool))
+    CONSOLE_ENDPOINT = ConfigEntry(YamlConfigEntry("console.endpoint"))
+    CA_CERT_FILE_PATH = ConfigEntry(YamlConfigEntry("admin.caCertFilePath"))
+    HTTP_PROXY_URL = ConfigEntry(YamlConfigEntry("admin.httpProxyURL"))
+
+
+class Credentials(object):
+    SECTION = "credentials"
+    COMMAND = ConfigEntry(YamlConfigEntry("admin.command", list))
+    """
+    This command is executed to return a token using an external process.
+    """
+
+    PROXY_COMMAND = ConfigEntry(YamlConfigEntry("admin.proxyCommand", list))
+    """
+    This command is executed to return a token for authorization with a proxy
+     in front of Flyte using an external process.
+    """
+
+    CLIENT_ID = ConfigEntry(YamlConfigEntry("admin.clientId"))
+    """
+    This is the public identifier for the app which handles authorization for a Flyte deployment.
+    More details here: https://www.oauth.com/oauth2-servers/client-registration/client-id-secret/.
+    """
+
+    CLIENT_CREDENTIALS_SECRET_LOCATION = ConfigEntry(YamlConfigEntry("admin.clientSecretLocation"))
+    """
+    Used for basic auth, which is automatically called during pyflyte. This will allow the Flyte engine to read the
+    password from a mounted file.
+    """
+
+    CLIENT_CREDENTIALS_SECRET_ENV_VAR = ConfigEntry(YamlConfigEntry("admin.clientSecretEnvVar"))
+    """
+    Used for basic auth, which is automatically called during pyflyte. This will allow the Flyte engine to read the
+    password from a mounted environment variable.
+    """
+
+    SCOPES = ConfigEntry(YamlConfigEntry("admin.scopes", list))
+    """
+    This setting can be used to manually pass in scopes into authenticator flows - eg.) for Auth0 compatibility
+    """
+
+    AUTH_MODE = ConfigEntry(YamlConfigEntry("admin.authType"))
+    """
+    The auth mode defines the behavior used to request and refresh credentials. The currently supported modes include:
+    - 'standard' or 'Pkce': This uses the pkce-enhanced authorization code flow by opening a browser window to initiate
+            credentials access.
+    - "DeviceFlow": This uses the Device Authorization Flow
+    - 'basic', 'client_credentials' or 'clientSecret': This uses symmetric key auth in which the end user enters a
+            client id and a client secret and public key encryption is used to facilitate authentication.
+    - None: No auth will be attempted.
+    """
+
+
+class Task(object):
+    ORG = ConfigEntry(YamlConfigEntry("task.org"))
+    PROJECT = ConfigEntry(YamlConfigEntry("task.project"))
+    DOMAIN = ConfigEntry(YamlConfigEntry("task.domain"))