flyte 2.0.0b32__py3-none-any.whl

flyte/_utils/__init__.py

@@ -0,0 +1,30 @@
+"""
+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, 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",
+    "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

@@ -0,0 +1,121 @@
+"""Manages an async event loop on another thread. Developers should only require to call
+sync to use the managed loop:
+
+from flytekit.tools.asyn import run_sync
+
+async def async_add(a: int, b: int) -> int:
+    return a + b
+
+result = run_sync(async_add, a=10, b=12)
+"""
+
+from __future__ import annotations
+
+import asyncio
+import atexit
+import functools
+import os
+import threading
+from contextlib import contextmanager
+from typing import Any, Awaitable, Callable, TypeVar
+
+from typing_extensions import ParamSpec
+
+from flyte._logging import logger
+
+T = TypeVar("T")
+
+P = ParamSpec("P")
+
+
+@contextmanager
+def _selector_policy():
+    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)
+
+
+class _TaskRunner:
+    """A task runner that runs an asyncio event loop on a background thread."""
+
+    def __init__(self) -> None:
+        self.__loop: asyncio.AbstractEventLoop | None = None
+        self.__runner_thread: threading.Thread | None = None
+        self.__lock = threading.Lock()
+        atexit.register(self._close)
+
+    def _close(self) -> None:
+        if self.__loop:
+            self.__loop.stop()
+
+    def _execute(self) -> None:
+        loop = self.__loop
+        assert loop is not None
+        try:
+            loop.run_forever()
+        finally:
+            loop.close()
+
+    def get_exc_handler(self):
+        def exc_handler(loop, context):
+            logger.error(
+                f"Taskrunner for {self.__runner_thread.name if self.__runner_thread else 'no thread'} caught"
+                f" exception in {loop}: {context}"
+            )
+
+        return exc_handler
+
+    def run(self, coro: Any) -> Any:
+        """Synchronously run a coroutine on a background thread."""
+        name = f"{threading.current_thread().name} : loop-runner"
+        with self.__lock:
+            if self.__loop is None:
+                with _selector_policy():
+                    self.__loop = asyncio.new_event_loop()
+
+                exc_handler = self.get_exc_handler()
+                self.__loop.set_exception_handler(exc_handler)
+                self.__runner_thread = threading.Thread(target=self._execute, daemon=True, name=name)
+                self.__runner_thread.start()
+        fut = asyncio.run_coroutine_threadsafe(coro, self.__loop)
+
+        res = fut.result(None)
+
+        return res
+
+
+class _AsyncLoopManager:
+    def __init__(self: _AsyncLoopManager):
+        self._runner_map: dict[str, _TaskRunner] = {}
+
+    def run_sync(self, coro_func: Callable[..., Awaitable[T]], *args, **kwargs) -> T:
+        """
+        This should be called from synchronous functions to run an async function.
+        """
+        name = threading.current_thread().name + f"PID:{os.getpid()}"
+        coro = coro_func(*args, **kwargs)
+        if name not in self._runner_map:
+            if len(self._runner_map) > 500:
+                logger.warning(
+                    "More than 500 event loop runners created!!! This could be a case of runaway recursion..."
+                )
+            self._runner_map[name] = _TaskRunner()
+        return self._runner_map[name].run(coro)
+
+    def synced(self, coro_func: Callable[P, Awaitable[T]]) -> Callable[P, T]:
+        """Make loop run coroutine until it returns. Runs in other thread"""
+
+        @functools.wraps(coro_func)
+        def wrapped(*args: Any, **kwargs: Any) -> T:
+            return self.run_sync(coro_func, *args, **kwargs)
+
+        return wrapped
+
+
+loop_manager = _AsyncLoopManager()
+run_sync = loop_manager.run_sync

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

@@ -0,0 +1,27 @@
+import asyncio
+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.
+
+    :param coros:
+    :param return_when:
+    :return:
+    """
+    # 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,
+    # so this is not strictly necessary.
+
+    for t in pending:  # type: asyncio.Task
+        t.cancel()  # Cancel all tasks that didn't finish first
+
+    for t in done:
+        err = t.exception()
+        if err:
+            raise err

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

@@ -0,0 +1,72 @@
+from __future__ import annotations
+
+import hashlib
+import os
+import pathlib
+import stat
+import typing
+from pathlib import Path
+from typing import List, Optional, Union
+
+from flyte._logging import logger
+
+
+def filehash_update(path: pathlib.Path, hasher: hashlib._Hash) -> None:
+    blocksize = 65536
+    with open(path, "rb") as f:
+        bytes = f.read(blocksize)
+        while bytes:
+            hasher.update(bytes)
+            bytes = f.read(blocksize)
+
+
+def _pathhash_update(path: Union[os.PathLike, str], hasher: hashlib._Hash) -> None:
+    path_list = str(path).split(os.sep)
+    hasher.update("".join(path_list).encode("utf-8"))
+
+
+def update_hasher_for_source(
+    source: Union[os.PathLike, List[os.PathLike]], hasher: hashlib._Hash, filter: Optional[typing.Callable] = None
+):
+    """
+    Walks the entirety of the source dir to compute a deterministic md5 hex digest of the dir contents.
+    :param os.PathLike source:
+    :param callable filter:
+    :return Text:
+    """
+
+    def compute_digest_for_file(path: os.PathLike, rel_path: os.PathLike) -> None:
+        # Only consider files that exist (e.g. disregard symlinks that point to non-existent files)
+        if not os.path.exists(path):
+            logger.info(f"Skipping non-existent file {path}")
+            return
+
+        # Skip socket files
+        if stat.S_ISSOCK(os.stat(path).st_mode):
+            logger.info(f"Skip socket file {path}")
+            return
+
+        if filter:
+            if filter(rel_path):
+                return
+
+        filehash_update(Path(path), hasher)
+        _pathhash_update(rel_path, hasher)
+
+    def compute_digest_for_dir(source: os.PathLike):
+        for root, _, files in os.walk(str(source), topdown=True):
+            files.sort()
+
+            for fname in files:
+                abspath = os.path.join(root, fname)
+                relpath = os.path.relpath(abspath, source)
+                compute_digest_for_file(Path(abspath), Path(relpath))
+
+    if isinstance(source, list):
+        for src in source:
+            if os.path.isdir(src):
+                compute_digest_for_dir(src)
+            else:
+                compute_digest_for_file(src, os.path.basename(src))
+    else:
+        compute_digest_for_dir(source)

flyte/_utils/helpers.py

@@ -0,0 +1,134 @@
+import os
+import string
+import typing
+from contextlib import contextmanager
+from pathlib import Path
+
+
+def load_proto_from_file(pb2_type, path):
+    with open(path, "rb") as reader:
+        out = pb2_type()
+        out.ParseFromString(reader.read())
+        return out
+
+
+def write_proto_to_file(proto, path):
+    Path(os.path.dirname(path)).mkdir(parents=True, exist_ok=True)
+    with open(path, "wb") as writer:
+        writer.write(proto.SerializeToString())
+
+
+def str2bool(value: typing.Optional[str]) -> bool:
+    """
+    Convert a string to a boolean. This is useful for parsing environment variables.
+    :param value: The string to convert to a boolean
+    :return: the boolean value
+    """
+    if value is None:
+        return False
+    return value.lower() in ("true", "t", "1")
+
+
+BASE36_ALPHABET = string.digits + string.ascii_lowercase  # 0-9 + a-z (36 characters)
+
+
+def base36_encode(byte_data: bytes) -> str:
+    """
+    This function expects to encode bytes coming from an hd5 hash function into a base36 encoded string.
+    md5 shas are limited to 128 bits, so the maximum byte value should easily fit into a 30 character long string.
+    If the input is too large howeer
+    """
+    # Convert bytes to a big integer
+    num = int.from_bytes(byte_data, byteorder="big")
+
+    # Convert integer to base36 string
+    if num == 0:
+        return BASE36_ALPHABET[0]
+
+    base36 = []
+    while num:
+        num, rem = divmod(num, 36)
+        base36.append(BASE36_ALPHABET[rem])
+    return "".join(reversed(base36))
+
+
+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,
+    but the intention is that it returns the source folder for an editable install if the current working directory
+    is inside the editable install project - if the code is inside an src/ folder, and the cwd is a level above,
+    it should still work, returning the src/ folder. If cwd is the src/ folder, this should return the same.
+
+    The idea is that the return path will be used to determine the relative path for imported modules when building
+    the code bundle.
+
+    :return:
+    """
+
+    from flyte._logging import logger
+
+    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
+    # the pyproject.toml file.
+    # Two scenarios to consider
+    #   - if cwd is nested inside the editable install folder.
+    #   - if the cwd is exactly one level above the editable install folder.
+    cwd = Path.cwd()
+    for install in editable_installs:
+        # child.is_relative_to(parent) is True if child is inside parent
+        if cwd.is_relative_to(install):
+            return install
+        else:
+            # check if the cwd is one level above the install folder
+            if install.parent == cwd:
+                # check if the install folder contains a pyproject.toml file
+                if (cwd / "pyproject.toml").exists() or (cwd / "setup.py").exists():
+                    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)