flyte 2.0.0b13__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/coro_management.py

@@ -11,7 +11,8 @@ async def run_coros(*coros: typing.Coroutine, return_when: str = asyncio.FIRST_C
     :param return_when:
     :return:
     """
-    tasks: typing.List[asyncio.Task[typing.Never]] = [asyncio.create_task(c) for c in coros]
+    # tasks: typing.List[asyncio.Task[typing.Never]] = [asyncio.create_task(c) for c in coros] # Python 3.11+
+    tasks: typing.List[asyncio.Task] = [asyncio.create_task(c) for c in coros]
     done, pending = await asyncio.wait(tasks, return_when=return_when)
     # TODO we might want to handle asyncio.CancelledError here, for cases when the `action` is cancelled
     # and we want to propagate it to all tasks. Though the backend will handle it anyway,

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.0b13'
-__version_tuple__ = version_tuple = (2, 0, 0, 'b13')
+__version__ = version = '2.0.0b30'
+__version_tuple__ = version_tuple = (2, 0, 0, 'b30')
 
-__commit_id__ = commit_id = 'g07f30e36d'
+__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

@@ -3,8 +3,7 @@ from pathlib import Path
 from types import ModuleType
 from typing import Any, Dict, List, cast
 
-import click
-from click import Context
+import rich_click as click
 
 import flyte
 
@@ -44,10 +43,8 @@ class BuildEnvCommand(click.Command):
         self.build_args = build_args
         super().__init__(*args, **kwargs)
 
-    def invoke(self, ctx: Context):
-        from rich.console import Console
-
-        console = Console()
+    def invoke(self, ctx: click.Context):
+        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):
+    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)
+        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