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

flyte/_tools.py

@@ -0,0 +1,27 @@
+import os
+
+
+def ipython_check() -> bool:
+    """
+    Check if interface is launching from iPython (not colab)
+    :return is_ipython (bool): True or False
+    """
+    is_ipython = False
+    try:  # Check if running interactively using ipython.
+        from IPython import get_ipython
+
+        if get_ipython() is not None:
+            is_ipython = True
+    except (ImportError, NameError):
+        pass
+    return is_ipython
+
+
+def is_in_cluster() -> bool:
+    """
+    Check if the task is running in a cluster
+    :return is_in_cluster (bool): True or False
+    """
+    if os.getenv("_UN_CLS"):
+        return True
+    return False

flyte/_trace.py

@@ -0,0 +1,120 @@
+import functools
+import inspect
+import time
+from datetime import timedelta
+from typing import Any, AsyncGenerator, AsyncIterator, Awaitable, Callable, TypeGuard, TypeVar, Union, cast
+
+from flyte.models import NativeInterface
+
+T = TypeVar("T")
+
+
+def trace(func: Callable[..., T]) -> Callable[..., T]:
+    """
+    A decorator that traces function execution with timing information.
+    Works with regular functions, async functions, and async generators/iterators.
+    """
+
+    @functools.wraps(func)
+    def wrapper_sync(*args: Any, **kwargs: Any) -> Any:
+        raise NotImplementedError
+
+    @functools.wraps(func)
+    async def wrapper_async(*args: Any, **kwargs: Any) -> Any:
+        from flyte._context import internal_ctx
+
+        ctx = internal_ctx()
+        if ctx.is_task_context():
+            # If we are in a task context, that implies we are executing a Run.
+            # In this scenario, we should submit the task to the controller.
+            # We will also check if we are not initialized, It is not expected to be not initialized
+            from ._internal.controllers import get_controller
+
+            controller = get_controller()
+            iface = NativeInterface.from_callable(func)
+            info, ok = await controller.get_action_outputs(iface, func, *args, **kwargs)
+            if ok:
+                if info.output:
+                    return info.output
+                elif info.error:
+                    raise info.error
+            start_time = time.time()
+            try:
+                # Cast to Awaitable to satisfy mypy
+                coroutine_result = cast(Awaitable[Any], func(*args, **kwargs))
+                results = await coroutine_result
+                duration = time.time() - start_time
+                info.add_outputs(results, timedelta(seconds=duration))
+                await controller.record_trace(info)
+                return results
+            except Exception as e:
+                # If there is an error, we need to record it
+                duration = time.time() - start_time
+                info.add_error(e, timedelta(seconds=duration))
+                await controller.record_trace(info)
+                raise e
+        else:
+            # If we are not in a task context, we can just call the function normally
+            # Cast to Awaitable to satisfy mypy
+            coroutine_result = cast(Awaitable[Any], func(*args, **kwargs))
+            return await coroutine_result
+
+    def is_async_iterable(obj: Any) -> TypeGuard[Union[AsyncGenerator, AsyncIterator]]:
+        return hasattr(obj, "__aiter__")
+
+    @functools.wraps(func)
+    async def wrapper_async_iterator(*args: Any, **kwargs: Any) -> AsyncIterator[Any]:
+        from flyte._context import internal_ctx
+
+        ctx = internal_ctx()
+        if ctx.is_task_context():
+            # If we are in a task context, that implies we are executing a Run.
+            # In this scenario, we should submit the task to the controller.
+            # We will also check if we are not initialized, It is not expected to be not initialized
+            from ._internal.controllers import get_controller
+
+            controller = get_controller()
+            iface = NativeInterface.from_callable(func)
+            info, ok = await controller.get_action_outputs(iface, func, *args, **kwargs)
+            if ok:
+                if info.output:
+                    for item in info.output:
+                        yield item
+                elif info.error:
+                    raise info.error
+            start_time = time.time()
+            try:
+                items = []
+                result = func(*args, **kwargs)
+                # TODO ideally we should use streaming into the type-engine so that it stream uploads large blocks
+                if inspect.isasyncgen(result) or is_async_iterable(result):
+                    # If it's directly an async generator
+                    async_iter = result
+                    async for item in async_iter:
+                        items.append(item)
+                        yield item
+                duration = time.time() - start_time
+                info.add_outputs(items, timedelta(seconds=duration))
+                await controller.record_trace(info)
+                return
+            except Exception as e:
+                end_time = time.time()
+                duration = end_time - start_time
+                info.add_error(e, timedelta(seconds=duration))
+                await controller.record_trace(info)
+                raise e
+        else:
+            result = func(*args, **kwargs)
+            if is_async_iterable(result):
+                async for item in result:
+                    yield item
+
+    # Choose the appropriate wrapper based on the function type
+    if inspect.iscoroutinefunction(func):
+        # This handles async functions that return normal values
+        return cast(Callable[..., T], wrapper_async)
+    elif inspect.isasyncgenfunction(func):
+        return cast(Callable[..., T], wrapper_async_iterator)
+    else:
+        # For regular sync functions
+        return cast(Callable[..., T], wrapper_sync)

flyte/_utils/__init__.py

@@ -0,0 +1,26 @@
+"""
+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 .lazy_module import lazy_module
+from .org_discovery import hostname_from_url, org_from_endpoint, sanitize_endpoint
+from .uv_script_parser import parse_uv_script_file
+
+__all__ = [
+    "AsyncLRUCache",
+    "filehash_update",
+    "get_cwd_editable_install",
+    "hostname_from_url",
+    "lazy_module",
+    "org_from_endpoint",
+    "parse_uv_script_file",
+    "run_coros",
+    "sanitize_endpoint",
+    "update_hasher_for_source",
+]

flyte/_utils/asyn.py

@@ -0,0 +1,119 @@
+"""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)
+"""
+
+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):
+        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,23 @@
+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]
+    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
+
+    for t in done:
+        err = t.exception()
+        if err:
+            raise err

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)