flyte 0.0.1b0__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,128 @@
+import functools
+import inspect
+import time
+from datetime import timedelta
+from typing import Any, AsyncGenerator, AsyncIterator, Awaitable, Callable, TypeGuard, TypeVar, Union, cast
+
+from flyte._datastructures 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.
+    """
+    func_name = func.__name__
+
+    @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 = await get_controller()
+            iface = NativeInterface.from_callable(func)
+            info, ok = await controller.get_action_outputs(iface, func_name, *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 = await get_controller()
+            iface = NativeInterface.from_callable(func)
+            info, ok = await controller.get_action_outputs(iface, func_name, *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):
+                    # If it's directly an async generator
+                    async_iter = result
+                    async for item in async_iter:
+                        items.append(item)
+                        yield item
+                elif is_async_iterable(result):
+                    # If it's an async iterable (has __aiter__)
+                    async for item in result:
+                        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
+        print(f"Coroutine function {func.__name__}")
+        return cast(Callable[..., T], wrapper_async)
+    elif inspect.isasyncgenfunction(func):
+        print(f"Async generator function {func.__name__}")
+        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,20 @@
+"""
+Internal utility functions.
+
+Except for logging, modules in this package should not depend on any other part of the repo.
+"""
+
+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 .uv_script_parser import parse_uv_script_file
+
+__all__ = [
+    "filehash_update",
+    "get_cwd_editable_install",
+    "lazy_module",
+    "parse_uv_script_file",
+    "run_coros",
+    "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/coro_management.py

@@ -0,0 +1,25 @@
+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
+        else:
+            print(f"Task result: {t.result()}")

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,108 @@
+import os
+import string
+import typing
+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))
+
+
+# does not work at all in the setuptools case. see old flytekit editable installs
+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:
+    """
+    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}")
+
+    # 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

flyte/_utils/lazy_module.py

@@ -0,0 +1,54 @@
+import importlib.util
+import sys
+import types
+
+
+class _LazyModule(types.ModuleType):
+    """
+    `lazy_module` returns an instance of this class if the module is not found in the python environment.
+    """
+
+    def __init__(self, module_name: str):
+        super().__init__(module_name)
+        self._module_name = module_name
+
+    def __getattribute__(self, attr):
+        raise ImportError(f"Module {object.__getattribute__(self, '_module_name')} is not yet installed.")
+
+
+def is_imported(module_name):
+    """
+    This function is used to check if a module has been imported by the regular import.
+    Return false if module is lazy imported and not used yet.
+    """
+    return (
+        module_name in sys.modules
+        and object.__getattribute__(lazy_module(module_name), "__class__").__name__ != "_LazyModule"
+    )
+
+
+def lazy_module(fullname):
+    """
+    This function is used to lazily import modules.  It is used in the following way:
+    .. code-block:: python
+        from flytekit.lazy_import import lazy_module
+        sklearn = lazy_module("sklearn")
+        sklearn.svm.SVC()
+    :param Text fullname: The full name of the module to import
+    """
+    if fullname in sys.modules:
+        return sys.modules[fullname]
+    # https://docs.python.org/3/library/importlib.html#implementing-lazy-imports
+    spec = importlib.util.find_spec(fullname)
+    if spec is None or spec.loader is None:
+        # Return a lazy module if the module is not found in the python environment,
+        # so that we can raise a proper error when the user tries to access an attribute in the module.
+        # The reason to do this is because importlib.util.LazyLoader still requires
+        # the module to be installed even if you don't use it.
+        return _LazyModule(fullname)
+    loader = importlib.util.LazyLoader(spec.loader)
+    spec.loader = loader
+    module = importlib.util.module_from_spec(spec)
+    sys.modules[fullname] = module
+    loader.exec_module(module)
+    return module

flyte/_utils/uv_script_parser.py

@@ -0,0 +1,49 @@
+import pathlib
+import re
+from dataclasses import dataclass, field
+from typing import Dict, List, Optional
+
+import toml
+
+
+@dataclass
+class ToolUVConfig:
+    exclude_newer: Optional[str] = None
+
+
+@dataclass
+class UVScriptMetadata:
+    requires_python: Optional[str] = None
+    dependencies: List[str] = field(default_factory=list)
+    tool: Optional[Dict[str, ToolUVConfig]] = None
+
+
+def _extract_uv_metadata_block(text: str) -> str | None:
+    pattern = re.compile(r"# /// script\s*(.*?)# ///", re.DOTALL)
+    match = pattern.search(text)
+    if not match:
+        return None
+    lines = [line.lstrip("# ").rstrip() for line in match.group(1).splitlines()]
+    return "\n".join(lines)
+
+
+def parse_uv_script_file(path: pathlib.Path) -> UVScriptMetadata:
+    if not path.exists() or not path.is_file():
+        raise FileNotFoundError(f"File not found: {path}")
+
+    text = path.read_text(encoding="utf-8")
+    raw_header = _extract_uv_metadata_block(text)
+    if raw_header is None:
+        raise ValueError("No uv metadata block found")
+
+    try:
+        data = toml.loads(raw_header)
+    except toml.TomlDecodeError as e:
+        raise ValueError(f"Invalid TOML in metadata block: {e}")
+
+    tool_data = data.get("tool", {}).get("uv", {})
+    return UVScriptMetadata(
+        requires_python=data.get("requires-python"),
+        dependencies=data.get("dependencies", []),
+        tool={"uv": ToolUVConfig(exclude_newer=tool_data.get("exclude-newer"))} if tool_data else None,
+    )

flyte/_version.py

@@ -0,0 +1,21 @@
+# file generated by setuptools-scm
+# don't change, don't track in version control
+
+__all__ = ["__version__", "__version_tuple__", "version", "version_tuple"]
+
+TYPE_CHECKING = False
+if TYPE_CHECKING:
+    from typing import Tuple
+    from typing import Union
+
+    VERSION_TUPLE = Tuple[Union[int, str], ...]
+else:
+    VERSION_TUPLE = object
+
+version: str
+__version__: str
+__version_tuple__: VERSION_TUPLE
+version_tuple: VERSION_TUPLE
+
+__version__ = version = '0.0.1b0'
+__version_tuple__ = version_tuple = (0, 0, 1, 'b0')