flyte 0.0.1b0__py3-none-any.whl

union/_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 union._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

union/_utils/file_handling.py

@@ -0,0 +1,71 @@
+from __future__ import annotations
+
+import hashlib
+import os
+import pathlib
+import stat
+from pathlib import Path
+from typing import List, Optional, Union
+
+from union._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 = 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[callable] = None
+) -> str:
+    """
+    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) -> None:
+        for root, _, files in os.walk(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(abspath, 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)

union/_utils/helpers.py

@@ -0,0 +1,46 @@
+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")
+
+
+BASE62_ALPHABET = string.digits + string.ascii_letters  # 0-9 + A-Z + a-z (62 characters)
+
+
+def base62_encode(byte_data: bytes) -> str:
+    # Convert bytes to a big integer
+    num = int.from_bytes(byte_data, byteorder="big")
+
+    # Convert integer to base62 string
+    if num == 0:
+        return BASE62_ALPHABET[0]
+
+    base62 = []
+    while num:
+        num, rem = divmod(num, 62)
+        base62.append(BASE62_ALPHABET[rem])
+    return "".join(reversed(base62))

union/_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

union/_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,
+    )

union/_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.1.dev160+g7ce653a.d20250515'
+__version_tuple__ = version_tuple = (0, 1, 'dev160', 'g7ce653a.d20250515')

union/errors.py

@@ -0,0 +1,128 @@
+"""
+Exceptions raised by Union.
+
+These errors are raised when the underlying task execution fails, either because of a user error, system error or an
+unknown error.
+"""
+
+from typing import Literal
+
+ErrorKind = Literal["system", "unknown", "user"]
+
+
+class RuntimeError(RuntimeError):
+    """
+    Base class for all Union runtime errors. These errors are raised when the underlying task execution fails, either
+    because of a user error, system error or an unknown error.
+    """
+
+    def __init__(self, code: str, kind: ErrorKind, root_cause_message: str, worker: str | None = None):
+        super().__init__(root_cause_message)
+        self.code = code
+        self.kind = kind
+        self.worker = worker
+
+
+class InitializationError(RuntimeError):
+    """
+    This error is raised when the Union system is tried to access without being initialized.
+    """
+
+
+class RuntimeSystemError(RuntimeError):
+    """
+    This error is raised when the underlying task execution fails because of a system error. This could be a bug in the
+    Union system or a bug in the user's code.
+    """
+
+    def __init__(self, code: str, message: str, worker: str | None = None):
+        super().__init__(code, "system", message, worker)
+
+
+class UnionRpcError(RuntimeSystemError):
+    """
+    This error is raised when communication with the Union server fails.
+    """
+
+
+class RuntimeUserError(RuntimeError):
+    """
+    This error is raised when the underlying task execution fails because of an error in the user's code.
+    """
+
+    def __init__(self, code: str, message: str, worker: str | None = None):
+        super().__init__(code, "user", message, worker)
+
+
+class RuntimeUnknownError(RuntimeError):
+    """
+    This error is raised when the underlying task execution fails because of an unknown error.
+    """
+
+    def __init__(self, code: str, message: str, worker: str | None = None):
+        super().__init__(code, "unknown", message, worker)
+
+
+class OOMError(RuntimeUserError):
+    """
+    This error is raised when the underlying task execution fails because of an out-of-memory error.
+    """
+
+
+class TaskInterruptedError(RuntimeUserError):
+    """
+    This error is raised when the underlying task execution is interrupted.
+    """
+
+
+class PrimaryContainerNotFoundError(RuntimeUserError):
+    """
+    This error is raised when the primary container is not found.
+    """
+
+
+class TaskTimeoutError(RuntimeUserError):
+    """
+    This error is raised when the underlying task execution runs for longer than the specified timeout.
+    """
+
+
+class RetriesExhaustedError(RuntimeUserError):
+    """
+    This error is raised when the underlying task execution fails after all retries have been exhausted.
+    """
+
+
+class InvalidImageNameError(RuntimeUserError):
+    """
+    This error is raised when the image name is invalid.
+    """
+
+
+class ImagePullBackOffError(RuntimeUserError):
+    """
+    This error is raised when the image cannot be pulled.
+    """
+
+
+class CustomError(RuntimeUserError):
+    """
+    This error is raised when the user raises a custom error.
+    """
+
+    def __init__(self, code: str, message: str):
+        super().__init__(code, message, "user")
+
+    @classmethod
+    def from_exception(cls, e: Exception):
+        """
+        Create a CustomError from an exception. The exception's class name is used as the error code and the exception
+        message is used as the error message.
+        """
+        return cls(e.__class__.__name__, str(e))
+
+
+class NotInTaskContextError(RuntimeUserError):
+    """
+    This error is raised when the user tries to access the task context outside of a task.
+    """

union/extras/__init__.py

@@ -0,0 +1,5 @@
+from ._container import ContainerTask
+
+__all__ = [
+    "ContainerTask",
+]