flyte 2.0.0b22__py3-none-any.whl → 2.0.0b30__py3-none-any.whl

flyte/_utils/__init__.py

@@ -9,12 +9,13 @@ from .coro_management import run_coros
 from .file_handling import filehash_update, update_hasher_for_source
 from .helpers import get_cwd_editable_install
 from .lazy_module import lazy_module
-from .module_loader import load_python_modules
+from .module_loader import adjust_sys_path, load_python_modules
 from .org_discovery import hostname_from_url, org_from_endpoint, sanitize_endpoint
 from .uv_script_parser import parse_uv_script_file
 
 __all__ = [
     "AsyncLRUCache",
+    "adjust_sys_path",
     "filehash_update",
     "get_cwd_editable_install",
     "hostname_from_url",

flyte/_utils/asyn.py

@@ -9,6 +9,8 @@ async def async_add(a: int, b: int) -> int:
 result = run_sync(async_add, a=10, b=12)
 """
 
+from __future__ import annotations
+
 import asyncio
 import atexit
 import functools
@@ -88,7 +90,7 @@ class _TaskRunner:
 
 
 class _AsyncLoopManager:
-    def __init__(self):
+    def __init__(self: _AsyncLoopManager):
         self._runner_map: dict[str, _TaskRunner] = {}
 
     def run_sync(self, coro_func: Callable[..., Awaitable[T]], *args, **kwargs) -> T:

flyte/_utils/docker_credentials.py

@@ -0,0 +1,173 @@
+"""Helper functions for creating Docker registry credentials for image pull secrets."""
+
+import base64
+import json
+import logging
+import os
+import subprocess
+from pathlib import Path
+from typing import Any
+
+logger = logging.getLogger(__name__)
+
+_CONFIG_JSON = "config.json"
+_DEFAULT_CONFIG_PATH = f"~/.docker/{_CONFIG_JSON}"
+_CRED_HELPERS = "credHelpers"
+_CREDS_STORE = "credsStore"
+
+
+def _load_docker_config(config_path: str | Path | None = None) -> dict[str, Any]:
+    """
+    Load Docker config from specified path.
+
+    Args:
+        config_path: Path to Docker config file. If None, uses DOCKER_CONFIG env var
+                    or defaults to ~/.docker/config.json
+
+    Returns:
+        Dictionary containing Docker config
+
+    Raises:
+        FileNotFoundError: If the config file does not exist
+        json.JSONDecodeError: If the config file is not valid JSON
+    """
+    if not config_path:
+        docker_config_env = os.environ.get("DOCKER_CONFIG")
+        if docker_config_env:
+            config_path = Path(docker_config_env) / _CONFIG_JSON
+        else:
+            config_path = Path(_DEFAULT_CONFIG_PATH).expanduser()
+    else:
+        config_path = Path(config_path).expanduser()
+
+    with open(config_path) as f:
+        return json.load(f)
+
+
+def _get_credential_helper(config: dict[str, Any], registry: str | None = None) -> str | None:
+    """Get credential helper for registry or global default."""
+    if registry and _CRED_HELPERS in config and registry in config[_CRED_HELPERS]:
+        return config[_CRED_HELPERS].get(registry)
+    return config.get(_CREDS_STORE)
+
+
+def _get_credentials_from_helper(helper: str, registry: str) -> tuple[str, str] | None:
+    """
+    Get credentials from system credential helper.
+
+    Args:
+        helper: Name of the credential helper (e.g., "osxkeychain", "wincred")
+        registry: Registry hostname to get credentials for
+
+    Returns:
+        Tuple of (username, password) or None if credentials cannot be retrieved
+    """
+    helper_cmd = f"docker-credential-{helper}"
+
+    try:
+        process = subprocess.Popen(
+            [helper_cmd, "get"],
+            stdin=subprocess.PIPE,
+            stdout=subprocess.PIPE,
+            stderr=subprocess.PIPE,
+            text=True,
+        )
+        output, error = process.communicate(input=registry)
+
+        if process.returncode != 0:
+            logger.error(f"Credential helper error: {error}")
+            return None
+
+        creds = json.loads(output)
+        return creds.get("Username"), creds.get("Secret")
+    except FileNotFoundError:
+        logger.error(f"Credential helper {helper_cmd} not found in PATH")
+        return None
+    except Exception as e:
+        logger.error(f"Error getting credentials: {e!s}")
+        return None
+
+
+def create_dockerconfigjson_from_config(
+    registries: list[str] | None = None,
+    docker_config_path: str | Path | None = None,
+) -> str:
+    """
+    Create a dockerconfigjson string from existing Docker config.
+
+    This function extracts Docker registry credentials from the user's Docker config file
+    and creates a JSON string containing only the credentials for the specified registries.
+    It handles credentials stored directly in the config file as well as those managed by
+    credential helpers.
+
+    Args:
+        registries: List of registries to extract credentials for. If None, all registries
+                   from the config will be used.
+        docker_config_path: Path to the Docker config file. If None, the function will look
+                          for the config file in the standard locations.
+
+    Returns:
+        JSON string in dockerconfigjson format: {"auths": {"registry": {"auth": "..."}}}
+
+    Raises:
+        FileNotFoundError: If Docker config file cannot be found
+        ValueError: If no credentials can be extracted
+    """
+    config = _load_docker_config(docker_config_path)
+
+    # Create new config structure with empty auths
+    new_config: dict[str, Any] = {"auths": {}}
+
+    # Use specified registries or all from config
+    target_registries = registries or list(config.get("auths", {}).keys())
+
+    if not target_registries:
+        raise ValueError("No registries found in Docker config and none specified")
+
+    for registry in target_registries:
+        registry_config = config.get("auths", {}).get(registry, {})
+        if registry_config.get("auth"):
+            # Direct auth token exists
+            new_config["auths"][registry] = {"auth": registry_config["auth"]}
+        else:
+            # Try to get credentials from helper
+            helper = _get_credential_helper(config, registry)
+            if helper:
+                creds = _get_credentials_from_helper(helper, registry)
+                if creds:
+                    username, password = creds
+                    auth_string = f"{username}:{password}"
+                    new_config["auths"][registry] = {"auth": base64.b64encode(auth_string.encode()).decode()}
+                else:
+                    logger.warning(f"Could not retrieve credentials for {registry} from credential helper")
+            else:
+                logger.warning(f"No credentials found for {registry}")
+
+    if not new_config["auths"]:
+        raise ValueError(f"No credentials could be extracted for registries: {', '.join(target_registries)}")
+
+    return json.dumps(new_config)
+
+
+def create_dockerconfigjson_from_credentials(
+    registry: str,
+    username: str,
+    password: str,
+) -> str:
+    """
+    Create a dockerconfigjson string from explicit credentials.
+
+    Args:
+        registry: Registry hostname (e.g., "ghcr.io", "docker.io")
+        username: Username or token name for the registry
+        password: Password or access token for the registry
+
+    Returns:
+        JSON string in dockerconfigjson format: {"auths": {"registry": {"auth": "..."}}}
+    """
+    auth_string = f"{username}:{password}"
+    auth_token = base64.b64encode(auth_string.encode()).decode()
+
+    config = {"auths": {registry: {"auth": auth_token}}}
+
+    return json.dumps(config)

flyte/_utils/module_loader.py

@@ -5,9 +5,9 @@ import sys
 from pathlib import Path
 from typing import List, Tuple
 
-from rich.progress import BarColumn, Progress, TextColumn, TimeElapsedColumn, TimeRemainingColumn
-
 import flyte.errors
+from flyte._constants import FLYTE_SYS_PATH
+from flyte._logging import logger
 
 
 def load_python_modules(path: Path, recursive: bool = False) -> Tuple[List[str], List[Tuple[Path, str]]]:
@@ -18,6 +18,8 @@ def load_python_modules(path: Path, recursive: bool = False) -> Tuple[List[str],
     :param recursive: If True, load modules recursively from subdirectories
     :return: List of loaded module names, and list of file paths that failed to load
     """
+    from rich.progress import BarColumn, Progress, TextColumn, TimeElapsedColumn, TimeRemainingColumn
+
     loaded_modules = []
     failed_paths = []
 
@@ -87,3 +89,16 @@ def _load_module_from_file(file_path: Path) -> str | None:
 
     except Exception as e:
         raise flyte.errors.ModuleLoadError(f"Failed to load module from {file_path}: {e}") from e
+
+
+def adjust_sys_path():
+    """
+    Adjust sys.path to include local sys.path entries under the root directory.
+    """
+    if "." not in sys.path or os.getcwd() not in sys.path:
+        sys.path.insert(0, ".")
+        logger.info(f"Added {os.getcwd()} to sys.path")
+    for p in os.environ.get(FLYTE_SYS_PATH, "").split(":"):
+        if p and p not in sys.path:
+            sys.path.insert(0, p)
+            logger.info(f"Added {p} to sys.path")

flyte/_version.py

@@ -28,7 +28,7 @@ version_tuple: VERSION_TUPLE
 commit_id: COMMIT_ID
 __commit_id__: COMMIT_ID
 
-__version__ = version = '2.0.0b22'
-__version_tuple__ = version_tuple = (2, 0, 0, 'b22')
+__version__ = version = '2.0.0b30'
+__version_tuple__ = version_tuple = (2, 0, 0, 'b30')
 
-__commit_id__ = commit_id = 'gce9a6bede'
+__commit_id__ = commit_id = 'gf1b5e2b2b'

flyte/cli/_abort.py

@@ -1,5 +1,4 @@
 import rich_click as click
-from rich.console import Console
 
 from flyte.cli import _common as common
 
@@ -23,6 +22,7 @@ def run(cfg: common.CLIConfig, run_name: str, project: str | None = None, domain
     cfg.init(project=project, domain=domain)
     r = Run.get(name=run_name)
     if r:
-        console = Console()
-        r.abort()
+        console = common.get_console()
+        with console.status(f"Aborting run '{run_name}'...", spinner="dots"):
+            r.abort()
         console.print(f"Run '{run_name}' has been aborted.")

flyte/cli/_build.py

@@ -44,9 +44,7 @@ class BuildEnvCommand(click.Command):
         super().__init__(*args, **kwargs)
 
     def invoke(self, ctx: click.Context):
-        from rich.console import Console
-
-        console = Console()
+        console = common.get_console()
         console.print(f"Building Environment: {self.obj_name}")
         obj: CLIConfig = ctx.obj
         obj.init()

flyte/cli/_common.py

@@ -1,11 +1,14 @@
 from __future__ import annotations
 
 import importlib.util
+import json
 import logging
 import os
+import pathlib
 import sys
 from abc import abstractmethod
 from dataclasses import dataclass, replace
+from functools import lru_cache
 from pathlib import Path
 from types import MappingProxyType, ModuleType
 from typing import Any, Dict, Iterable, List, Literal, Optional
@@ -19,10 +22,12 @@ from rich.pretty import pretty_repr
 from rich.table import Table
 from rich.traceback import Traceback
 
+import flyte.config
 import flyte.errors
+from flyte._logging import LogFormat
 from flyte.config import Config
 
-OutputFormat = Literal["table", "json", "table-simple"]
+OutputFormat = Literal["table", "json", "table-simple", "json-raw"]
 
 PREFERRED_BORDER_COLOR = "dim cyan"
 PREFERRED_ACCENT_COLOR = "bold #FFD700"
@@ -99,6 +104,7 @@ class CLIConfig:
     config: Config
     ctx: click.Context
     log_level: int | None = logging.ERROR
+    log_format: LogFormat = "console"
     endpoint: str | None = None
     insecure: bool = False
     org: str | None = None
@@ -111,13 +117,20 @@ class CLIConfig:
         """
         return replace(self, **kwargs)
 
-    def init(self, project: str | None = None, domain: str | None = None, root_dir: str | None = None):
+    def init(
+        self,
+        project: str | None = None,
+        domain: str | None = None,
+        root_dir: str | None = None,
+        images: tuple[str, ...] | None = None,
+        sync_local_sys_paths: bool = True,
+    ):
         from flyte.config._config import TaskConfig
 
         task_cfg = TaskConfig(
             org=self.org or self.config.task.org,
-            project=project or self.config.task.project,
-            domain=domain or self.config.task.domain,
+            project=project if project is not None else self.config.task.project,
+            domain=domain if domain is not None else self.config.task.domain,
         )
 
         kwargs: Dict[str, Any] = {}
@@ -131,7 +144,14 @@ class CLIConfig:
 
         updated_config = self.config.with_params(platform_cfg, task_cfg)
 
-        flyte.init_from_config(updated_config, log_level=self.log_level, root_dir=root_dir)
+        flyte.init_from_config(
+            updated_config,
+            log_level=self.log_level,
+            log_format=self.log_format,
+            root_dir=pathlib.Path(root_dir) if root_dir else None,
+            images=images,
+            sync_local_sys_paths=sync_local_sys_paths,
+        )
 
 
 class InvokeBaseMixin:
@@ -177,7 +197,7 @@ class InvokeBaseMixin:
         except Exception as e:
             if ctx.obj and ctx.obj.log_level and ctx.obj.log_level <= logging.DEBUG:
                 # If the user has requested verbose output, print the full traceback
-                console = Console()
+                console = get_console()
                 console.print(Traceback.from_exception(type(e), e, e.__traceback__))
                 exit(1)
             else:
@@ -354,7 +374,7 @@ def _table_format(table: Table, vals: Iterable[Any]) -> Table:
         if headers is None:
             headers = [k for k, _ in o]
             for h in headers:
-                table.add_column(h.capitalize())
+                table.add_column(h.capitalize(), no_wrap=True if "name" in h.casefold() else False)
         table.add_row(*[str(v) for _, v in o])
     return table
 
@@ -374,6 +394,7 @@ def format(title: str, vals: Iterable[Any], of: OutputFormat = "table") -> Table
                     header_style=HEADER_STYLE,
                     show_header=True,
                     border_style=PREFERRED_BORDER_COLOR,
+                    expand=True,
                 ),
                 vals,
             )
@@ -381,6 +402,11 @@ def format(title: str, vals: Iterable[Any], of: OutputFormat = "table") -> Table
             if not vals:
                 return pretty_repr([])
             return pretty_repr([v.to_dict() for v in vals])
+        case "json-raw":
+            if not vals:
+                return []
+            return json.dumps([v.to_dict() for v in vals])
+
     raise click.ClickException("Unknown output format. Supported formats are: table, table-simple, json.")
 
 
@@ -395,3 +421,48 @@ def get_panel(title: str, renderable: Any, of: OutputFormat = "table") -> Panel:
         title=f"[{PREFERRED_ACCENT_COLOR}]{title}[/{PREFERRED_ACCENT_COLOR}]",
         border_style=PREFERRED_BORDER_COLOR,
     )
+
+
+def get_console() -> Console:
+    """
+    Get a console that is configured to use colors if the terminal supports it.
+    """
+    return Console(color_system="auto", force_terminal=True, width=120)
+
+
+def parse_images(cfg: Config, values: tuple[str, ...] | None) -> None:
+    """
+    Parse image values and update the config.
+
+    Args:
+        cfg: The Config object to write images to
+        values: List of image strings in format "imagename=imageuri" or just "imageuri"
+    """
+    if values is None:
+        return
+    for value in values:
+        if "=" in value:
+            image_name, image_uri = value.split("=", 1)
+            cfg.image.image_refs[image_name] = image_uri
+        else:
+            # If no name specified, use "default" as the name
+            cfg.image.image_refs["default"] = value
+
+
+@lru_cache()
+def initialize_config(
+    ctx: click.Context,
+    project: str,
+    domain: str,
+    root_dir: str | None = None,
+    images: tuple[str, ...] | None = None,
+    sync_local_sys_paths: bool = True,
+):
+    obj: CLIConfig | None = ctx.obj
+    if obj is None:
+        import flyte.config
+
+        obj = CLIConfig(flyte.config.auto(), ctx)
+
+    obj.init(project, domain, root_dir, images, sync_local_sys_paths)
+    return obj

flyte/cli/_create.py

@@ -3,6 +3,7 @@ from typing import Any, Dict, get_args
 
 import rich_click as click
 
+import flyte
 import flyte.cli._common as common
 from flyte.cli._option import MutuallyExclusiveOption
 from flyte.remote import SecretTypes
@@ -23,18 +24,49 @@ def create():
     prompt="Enter secret value",
     hide_input=True,
     cls=MutuallyExclusiveOption,
-    mutually_exclusive=["from_file"],
+    mutually_exclusive=["from_file", "from_docker_config", "registry"],
 )
 @click.option(
     "--from-file",
     type=click.Path(exists=True),
     help="Path to the file with the binary secret.",
     cls=MutuallyExclusiveOption,
-    mutually_exclusive=["value"],
+    mutually_exclusive=["value", "from_docker_config", "registry"],
 )
 @click.option(
     "--type", type=click.Choice(get_args(SecretTypes)), default="regular", help="Type of the secret.", show_default=True
 )
+@click.option(
+    "--from-docker-config",
+    is_flag=True,
+    help="Create image pull secret from Docker config file (only for --type image_pull).",
+    cls=MutuallyExclusiveOption,
+    mutually_exclusive=["value", "from_file", "registry", "username", "password"],
+)
+@click.option(
+    "--docker-config-path",
+    type=click.Path(exists=True),
+    help="Path to Docker config file (defaults to ~/.docker/config.json or $DOCKER_CONFIG).",
+)
+@click.option(
+    "--registries",
+    help="Comma-separated list of registries to include (only with --from-docker-config).",
+)
+@click.option(
+    "--registry",
+    help="Registry hostname (e.g., ghcr.io, docker.io) for explicit credentials (only for --type image_pull).",
+    cls=MutuallyExclusiveOption,
+    mutually_exclusive=["value", "from_file", "from_docker_config"],
+)
+@click.option(
+    "--username",
+    help="Username for the registry (only with --registry).",
+)
+@click.option(
+    "--password",
+    help="Password for the registry (only with --registry). If not provided, will prompt.",
+    hide_input=True,
+)
 @click.pass_obj
 def secret(
     cfg: common.CLIConfig,
@@ -42,6 +74,12 @@ def secret(
     value: str | bytes | None = None,
     from_file: str | None = None,
     type: SecretTypes = "regular",
+    from_docker_config: bool = False,
+    docker_config_path: str | None = None,
+    registries: str | None = None,
+    registry: str | None = None,
+    username: str | None = None,
+    password: str | None = None,
     project: str | None = None,
     domain: str | None = None,
 ):
@@ -72,16 +110,80 @@ def secret(
     Other secrets should be specified as `regular`.
     If no type is specified, `regular` is assumed.
 
+    For image pull secrets, you have several options:
+
+    1. Interactive mode (prompts for registry, username, password):
     ```bash
     $ flyte create secret my_secret --type image_pull
     ```
+
+    2. With explicit credentials:
+    ```bash
+    $ flyte create secret my_secret --type image_pull --registry ghcr.io --username myuser
+    ```
+
+    3. Lastly, you can create a secret from your existing Docker installation (i.e., you've run `docker login` in
+    the past) and you just want to pull from those credentials. Since you may have logged in to multiple registries,
+    you can specify which registries to include. If no registries are specified, all registries are added.
+    ```bash
+    $ flyte create secret my_secret --type image_pull --from-docker-config --registries ghcr.io,docker.io
+    ```
     """
     from flyte.remote import Secret
 
+    # todo: remove this hack when secrets creation more easily distinguishes between org and project/domain level
+    #   (and domain level) secrets
+    project = "" if project is None else project
+    domain = "" if domain is None else domain
     cfg.init(project, domain)
-    if from_file:
+
+    # Handle image pull secret creation
+    if type == "image_pull":
+        if project != "" or domain != "":
+            raise click.ClickException("Project and domain must not be set when creating an image pull secret.")
+
+        if from_docker_config:
+            # Mode 3: From Docker config
+            from flyte._utils.docker_credentials import create_dockerconfigjson_from_config
+
+            registry_list = [r.strip() for r in registries.split(",")] if registries else None
+            try:
+                value = create_dockerconfigjson_from_config(
+                    registries=registry_list,
+                    docker_config_path=docker_config_path,
+                )
+            except Exception as e:
+                raise click.ClickException(f"Failed to create dockerconfigjson from Docker config: {e}") from e
+
+        elif registry:
+            # Mode 2: Explicit credentials
+            from flyte._utils.docker_credentials import create_dockerconfigjson_from_credentials
+
+            if not username:
+                username = click.prompt("Username")
+            if not password:
+                password = click.prompt("Password", hide_input=True)
+
+            value = create_dockerconfigjson_from_credentials(registry, username, password)
+
+        else:
+            # Mode 1: Interactive prompts
+            from flyte._utils.docker_credentials import create_dockerconfigjson_from_credentials
+
+            registry = click.prompt("Registry (e.g., ghcr.io, docker.io)")
+            username = click.prompt("Username")
+            password = click.prompt("Password", hide_input=True)
+
+            value = create_dockerconfigjson_from_credentials(registry, username, password)
+
+    elif from_file:
         with open(from_file, "rb") as f:
             value = f.read()
+
+    # Encode string values to bytes
+    if isinstance(value, str):
+        value = value.encode("utf-8")
+
     Secret.create(name=name, value=value, type=type)
 
 
@@ -194,3 +296,76 @@ def config(
         yaml.dump(d, f)
 
     click.echo(f"Config file written to {output_path}")
+
+
+@create.command(cls=common.CommandBase)
+@click.argument("task_name", type=str, required=True)
+@click.argument("name", type=str, required=True)
+@click.option(
+    "--schedule",
+    type=str,
+    required=True,
+    help="Cron schedule for the trigger. Defaults to every minute.",
+    show_default=True,
+)
+@click.option(
+    "--description",
+    type=str,
+    default="",
+    help="Description of the trigger.",
+    show_default=True,
+)
+@click.option(
+    "--auto-activate",
+    is_flag=True,
+    default=True,
+    help="Whether the trigger should not be automatically activated. Defaults to True.",
+    show_default=True,
+)
+@click.option(
+    "--trigger-time-var",
+    type=str,
+    default="trigger_time",
+    help="Variable name for the trigger time in the task inputs. Defaults to 'trigger_time'.",
+    show_default=True,
+)
+@click.pass_obj
+def trigger(
+    cfg: common.CLIConfig,
+    task_name: str,
+    name: str,
+    schedule: str,
+    trigger_time_var: str = "trigger_time",
+    auto_activate: bool = True,
+    description: str = "",
+    project: str | None = None,
+    domain: str | None = None,
+):
+    """
+    Create a new trigger for a task. The task name and trigger name are required.
+
+    Example:
+
+    ```bash
+    $ flyte create trigger my_task my_trigger --schedule "0 0 * * *"
+    ```
+
+    This will create a trigger that runs every day at midnight.
+    """
+    from flyte.remote import Trigger
+
+    cfg.init(project, domain)
+    console = common.get_console()
+
+    trigger = flyte.Trigger(
+        name=name,
+        automation=flyte.Cron(schedule),
+        description=description,
+        auto_activate=auto_activate,
+        inputs={trigger_time_var: flyte.TriggerTime},  # Use the trigger time variable in inputs
+        env_vars=None,
+        interruptible=None,
+    )
+    with console.status("Creating trigger..."):
+        v = Trigger.create(trigger, task_name=task_name)
+    console.print(f"[bold green]Trigger {v.name} created successfully![/bold green]")