flyte 0.0.1b0__py3-none-any.whl → 2.0.0b46__py3-none-any.whl

flyte/_utils/__init__.py

@@ -4,17 +4,28 @@ Internal utility functions.
 Except for logging, modules in this package should not depend on any other part of the repo.
 """
 
+from .async_cache import AsyncLRUCache
 from .coro_management import run_coros
 from .file_handling import filehash_update, update_hasher_for_source
-from .helpers import get_cwd_editable_install
+from .helpers import get_cwd_editable_install, str2bool
 from .lazy_module import lazy_module
+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",
+    "description_parser",
     "filehash_update",
     "get_cwd_editable_install",
+    "hostname_from_url",
     "lazy_module",
+    "load_python_modules",
+    "org_from_endpoint",
     "parse_uv_script_file",
     "run_coros",
+    "sanitize_endpoint",
+    "str2bool",
     "update_hasher_for_source",
 ]

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/async_cache.py

@@ -0,0 +1,139 @@
+import asyncio
+import time
+from collections import OrderedDict
+from typing import Awaitable, Callable, Dict, Generic, Optional, TypeVar
+
+K = TypeVar("K")
+V = TypeVar("V")
+
+
+class AsyncLRUCache(Generic[K, V]):
+    """
+    A high-performance async-compatible LRU cache.
+
+    Examples:
+    ```python
+    # Create a cache instance
+    cache = AsyncLRUCache[str, dict](maxsize=100)
+
+    async def fetch_data(user_id: str) -> dict:
+        # Define the expensive operation as a local function
+        async def get_user_data():
+            await asyncio.sleep(1)  # Simulating network/DB delay
+            return {"id": user_id, "name": f"User {user_id}"}
+
+    # Use the cache
+    return await cache.get(f"user:{user_id}", get_user_data)
+    ```
+    This cache can be used from async coroutines and handles concurrent access safely.
+    """
+
+    def __init__(self, maxsize: int = 128, ttl: Optional[float] = None):
+        """
+        Initialize the async LRU cache.
+
+        Args:
+            maxsize: Maximum number of items to keep in the cache
+            ttl: Time-to-live for cache entries in seconds, or None for no expiration
+        """
+        self._cache: OrderedDict[K, tuple[V, float]] = OrderedDict()
+        self._maxsize = maxsize
+        self._ttl = ttl
+        self._locks: Dict[K, asyncio.Lock] = {}
+        self._access_lock = asyncio.Lock()
+
+    async def get(self, key: K, value_func: Callable[[], V | Awaitable[V]]) -> V:
+        """
+        Get a value from the cache, computing it if necessary.
+
+        Args:
+            key: The cache key
+            value_func: Function or coroutine to compute the value if not cached
+
+        Returns:
+            The cached or computed value
+        """
+        # Fast path: check if key exists and is not expired
+        if key in self._cache:
+            value, timestamp = self._cache[key]
+            if self._ttl is None or time.time() - timestamp < self._ttl:
+                # Move the accessed item to the end (most recently used)
+                async with self._access_lock:
+                    self._cache.move_to_end(key)
+                return value
+
+        # Slow path: compute the value
+        # Get or create a lock for this key to prevent redundant computation
+        async with self._access_lock:
+            lock = self._locks.get(key)
+            if lock is None:
+                lock = asyncio.Lock()
+                self._locks[key] = lock
+
+        async with lock:
+            # Check again in case another coroutine computed the value while we waited
+            if key in self._cache:
+                value, timestamp = self._cache[key]
+                if self._ttl is None or time.time() - timestamp < self._ttl:
+                    async with self._access_lock:
+                        self._cache.move_to_end(key)
+                    return value
+
+            # Compute the value
+            if asyncio.iscoroutinefunction(value_func):
+                value = await value_func()
+            else:
+                value = value_func()  # type: ignore
+
+            # Store in cache
+            async with self._access_lock:
+                self._cache[key] = (value, time.time())
+                # Evict least recently used items if needed
+                while len(self._cache) > self._maxsize:
+                    self._cache.popitem(last=False)
+                # Clean up the lock
+                self._locks.pop(key, None)
+
+            return value
+
+    async def set(self, key: K, value: V) -> None:
+        """
+        Explicitly set a value in the cache.
+
+        Args:
+            key: The cache key
+            value: The value to cache
+        """
+        async with self._access_lock:
+            self._cache[key] = (value, time.time())
+            # Evict least recently used items if needed
+            while len(self._cache) > self._maxsize:
+                self._cache.popitem(last=False)
+
+    async def invalidate(self, key: K) -> None:
+        """Remove a specific key from the cache."""
+        async with self._access_lock:
+            self._cache.pop(key, None)
+
+    async def clear(self) -> None:
+        """Clear the entire cache."""
+        async with self._access_lock:
+            self._cache.clear()
+            self._locks.clear()
+
+    async def contains(self, key: K) -> bool:
+        """Check if a key exists in the cache and is not expired."""
+        if key not in self._cache:
+            return False
+
+        if self._ttl is None:
+            return True
+
+        _, timestamp = self._cache[key]
+        return time.time() - timestamp < self._ttl
+
+
+# Example usage:
+"""
+
+"""

flyte/_utils/coro_management.py

@@ -5,21 +5,22 @@ import typing
 async def run_coros(*coros: typing.Coroutine, return_when: str = asyncio.FIRST_COMPLETED):
     """
     Run a list of coroutines concurrently and wait for the first one to finish or exit.
-    When the first one finishes, cancel all other tasks.
+    When the first one finishes, cancel all other tasks. This helper function does not propagate CancelledError, but
+    will cancel pending tasks.
 
     :param coros:
     :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)
 
     for t in pending:  # type: asyncio.Task
         t.cancel()  # Cancel all tasks that didn't finish first
 
+    # Check for exceptions only in the completed tasks
     for t in done:
         err = t.exception()
         if err:
             raise err
-        else:
-            print(f"Task result: {t.result()}")

flyte/_utils/description_parser.py

@@ -0,0 +1,19 @@
+def parse_description(description: str, max_length: int = 255):
+    """
+    Parse and truncate a description string to fit within a maximum length.
+
+    If the description exceeds max_length, it will be truncated and suffixed with
+    "...(tr.)" to indicate truncation.
+
+    Args:
+        description: The description string to parse.
+        max_length: Maximum allowed length for the description. Defaults to 255.
+
+    Returns:
+        The parsed description string, truncated if necessary.
+    """
+    if len(description) <= max_length:
+        return description
+    if max_length >= 8:
+        return description[: max_length - 8] + "...(tr.)"
+    return description[:max_length]

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/helpers.py

@@ -1,6 +1,7 @@
 import os
 import string
 import typing
+from contextlib import contextmanager
 from pathlib import Path
 
 
@@ -51,7 +52,33 @@ def base36_encode(byte_data: bytes) -> str:
     return "".join(reversed(base36))
 
 
-# does not work at all in the setuptools case. see old flytekit editable installs
+def _iter_editable():
+    """
+    Yield (project_name, source_path) for every editable distribution
+    visible to the current interpreter
+    """
+    import json
+    import pathlib
+    from importlib.metadata import distributions
+
+    for dist in distributions():
+        # PEP-610 / PEP-660 (preferred, wheel-style editables)
+        direct = dist.read_text("direct_url.json")
+        if direct:
+            data = json.loads(direct)
+            if data.get("dir_info", {}).get("editable"):  # spec key
+                # todo: will need testing on windows
+                yield dist.metadata["Name"], pathlib.Path(data["url"][7:])  # strip file://
+                continue
+
+        # Legacy setuptools-develop / pip-e (egg-link)
+        for file in dist.files or ():  # importlib.metadata 3.8+
+            if file.suffix == ".egg-link":
+                with open(dist.locate_file(file), "r") as f:
+                    line = f.readline()
+                    yield dist.metadata["Name"], pathlib.Path(line.strip())
+
+
 def get_cwd_editable_install() -> typing.Optional[Path]:
     """
     This helper function is incomplete since it hasn't been tested with all the package managers out there,
@@ -64,28 +91,13 @@ def get_cwd_editable_install() -> typing.Optional[Path]:
 
     :return:
     """
-    import site
 
     from flyte._logging import logger
 
-    egg_links = [Path(p) for p in Path(site.getsitepackages()[0]).glob("*.egg-link")]
-    pth_files = [Path(p) for p in Path(site.getsitepackages()[0]).glob("*.pth")]
-
-    if not egg_links and not pth_files:
-        logger.debug("No editable installs found.")
-        return None
-
     editable_installs = []
-    egg_links.extend(pth_files)
-    for file in egg_links:
-        with open(file, "r") as f:
-            line = f.readline()
-            if line:
-                # Check if the first line is a directory
-                p = Path(line)
-                if p.is_dir():
-                    editable_installs.append(p)
-    logger.debug(f"Editable installs: {editable_installs}")
+    for name, path in _iter_editable():
+        logger.debug(f"Detected editable install: {name} at {path}")
+        editable_installs.append(path)
 
     # check to see if the current working directory is in any of the editable installs
     # including if the current folder is the root folder, one level up from the src and contains
@@ -106,3 +118,17 @@ def get_cwd_editable_install() -> typing.Optional[Path]:
                     return install  # note we want the install folder, not the parent
 
     return None
+
+
+@contextmanager
+def _selector_policy():
+    import asyncio
+
+    original_policy = asyncio.get_event_loop_policy()
+    try:
+        if os.name == "nt" and hasattr(asyncio, "WindowsSelectorEventLoopPolicy"):
+            asyncio.set_event_loop_policy(asyncio.WindowsSelectorEventLoopPolicy())
+
+        yield
+    finally:
+        asyncio.set_event_loop_policy(original_policy)

flyte/_utils/module_loader.py

@@ -0,0 +1,123 @@
+import importlib.util
+import os
+import sys
+from pathlib import Path
+from types import ModuleType
+from typing import List, Tuple
+
+import flyte.errors
+from flyte._constants import FLYTE_SYS_PATH
+from flyte._logging import logger
+
+
+def load_python_modules(
+    path: Path, root_dir: Path, recursive: bool = False
+) -> Tuple[List[ModuleType], List[Tuple[Path, str]]]:
+    """
+    Load all Python modules from a path and return list of loaded module names.
+
+    :param path: File or directory path
+    :param root_dir: Root directory to search for modules
+    :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 = []
+
+    if path.is_file() and path.suffix == ".py":
+        rel_path = path.resolve().relative_to(root_dir)
+        mod = (".".join(rel_path.parts))[:-3]
+        imported_module = importlib.import_module(mod)
+        loaded_modules.append(imported_module)
+
+    elif path.is_dir():
+        # Directory case - find all Python files
+        pattern = "**/*.py" if recursive else "*.py"
+        python_files = list(path.glob(pattern))
+
+        # Filter out __init__.py files
+        python_files = [f for f in python_files if f.name != "__init__.py"]
+
+        if not python_files:
+            # If no .py files found, try importing as a module
+            try:
+                rel_path = path.resolve().relative_to(root_dir)
+                mod = ".".join(rel_path.parts)
+                imported_module = importlib.import_module(mod)
+                loaded_modules.append(imported_module)
+            except (ValueError, ModuleNotFoundError):
+                pass
+        else:
+            with Progress(
+                TextColumn("[progress.description]{task.description}"),
+                BarColumn(),
+                "[progress.percentage]{task.percentage:>3.0f}%",
+                TimeElapsedColumn(),
+                TimeRemainingColumn(),
+                TextColumn("• {task.fields[current_file]}"),
+            ) as progress:
+                task = progress.add_task(f"Loading {len(python_files)} files", total=len(python_files), current_file="")
+
+                for file_path in python_files:
+                    progress.update(task, advance=1, current_file=file_path.name)
+
+                    try:
+                        rel_path = file_path.resolve().relative_to(root_dir)
+                        mod = (".".join(rel_path.parts))[:-3]
+                        imported_module = importlib.import_module(mod)
+                        loaded_modules.append(imported_module)
+                    except flyte.errors.ModuleLoadError as e:
+                        failed_paths.append((file_path, str(e)))
+
+                progress.update(task, current_file="[green]Done[/green]")
+
+    return loaded_modules, failed_paths
+
+
+def _load_module_from_file(file_path: Path) -> str | None:
+    """
+    Load a Python module from a file path.
+
+    :param file_path: Path to the Python file
+    :return: Module name if successfully loaded, None otherwise
+    """
+    try:
+        # Use the file stem as module name
+        module_name = file_path.stem
+
+        # Load the module specification
+        spec = importlib.util.spec_from_file_location(module_name, file_path)
+        if spec is None or spec.loader is None:
+            return None
+
+        # Create and execute the module
+        module = importlib.util.module_from_spec(spec)
+        sys.modules[module_name] = module
+        module_path = os.path.dirname(os.path.abspath(file_path))
+        sys.path.append(module_path)
+        spec.loader.exec_module(module)
+
+        return module_name
+
+    except Exception as e:
+        raise flyte.errors.ModuleLoadError(f"Failed to load module from {file_path}: {e}") from e
+
+
+def adjust_sys_path(additional_paths: List[str] | None = None):
+    """
+    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")
+    if additional_paths:
+        for p in additional_paths:
+            if p and p not in sys.path:
+                sys.path.insert(0, p)
+                logger.info(f"Added {p} to sys.path")

flyte/_utils/org_discovery.py

@@ -0,0 +1,57 @@
+def hostname_from_url(url: str) -> str:
+    """Parse a URL and return the hostname part."""
+
+    # Handle dns:/// format specifically (gRPC convention)
+    if url.startswith("dns:///"):
+        return url[7:]  # Skip the "dns:///" prefix
+
+    # Handle standard URL formats
+    import urllib.parse
+
+    parsed = urllib.parse.urlparse(url)
+    return parsed.netloc or parsed.path.lstrip("/").rsplit("/")[0]
+
+
+def org_from_endpoint(endpoint: str | None) -> str | None:
+    """
+    Extracts the organization from the endpoint URL. The organization is assumed to be the first part of the domain.
+    This is temporary until we have a proper organization discovery mechanism through APIs.
+
+    :param endpoint: The endpoint URL
+    :return: The organization name or None if not found
+    """
+    if not endpoint:
+        return None
+
+    hostname = hostname_from_url(endpoint)
+    domain_parts = hostname.split(".")
+    if len(domain_parts) > 2:
+        # Assuming the organization is the first part of the domain
+        return domain_parts[0]
+    return None
+
+
+def sanitize_endpoint(endpoint: str | None) -> str | None:
+    """
+    Sanitize the endpoint URL by ensuring it has a valid scheme.
+    :param endpoint: The endpoint URL to sanitize
+    :return: Sanitized endpoint URL or None if the input was None
+    """
+    if not endpoint:
+        return None
+    if "://" not in endpoint:
+        endpoint = f"dns:///{endpoint}"
+    else:
+        if endpoint.startswith("https://"):
+            # If the endpoint starts with dns:///, we assume it's a gRPC endpoint
+            endpoint = f"dns:///{endpoint[8:]}"
+        elif endpoint.startswith("http://"):
+            # If the endpoint starts with http://, we assume it's a REST endpoint
+            endpoint = f"dns:///{endpoint[7:]}"
+        elif not endpoint.startswith("dns:///"):
+            raise RuntimeError(
+                f"Invalid endpoint {endpoint}, expected format is "
+                f"dns:///<hostname> or https://<hostname> or http://<hostname>"
+            )
+    endpoint = endpoint.removesuffix("/")
+    return endpoint