flyte 2.0.0b32__py3-none-any.whl

flyte/_cache/cache.py

@@ -0,0 +1,147 @@
+import hashlib
+from dataclasses import dataclass, field
+from typing import (
+    Callable,
+    Generic,
+    List,
+    Optional,
+    Protocol,
+    Tuple,
+    Union,
+    runtime_checkable,
+)
+
+import rich.repr
+from typing_extensions import Literal, ParamSpec, TypeVar, get_args
+
+# if TYPE_CHECKING:
+from flyte._image import Image
+from flyte.models import CodeBundle
+
+P = ParamSpec("P")
+FuncOut = TypeVar("FuncOut")
+
+CacheBehavior = Literal["auto", "override", "disable"]
+
+
+@dataclass
+class VersionParameters(Generic[P, FuncOut]):
+    """
+    Parameters used for cache version hash generation.
+
+    :param func: The function to generate a version for. This is a required parameter but can be any callable
+    :type func: Callable[P, FuncOut]
+    :param image: The container image to generate a version for. This can be a string representing the
+        image name or an Image object.
+    :type image: Optional[Union[str, Image]]
+    """
+
+    func: Callable[P, FuncOut] | None
+    image: Optional[Union[str, Image]] = None
+    code_bundle: Optional[CodeBundle] = None
+
+
+@runtime_checkable
+class CachePolicy(Protocol):
+    def get_version(self, salt: str, params: VersionParameters) -> str: ...
+
+
+@rich.repr.auto
+@dataclass
+class Cache:
+    """
+    Cache configuration for a task.
+    :param behavior: The behavior of the cache. Can be "auto", "override" or "disable".
+    :param version_override: The version of the cache. If not provided, the version will be
+     generated based on the cache policies
+    :type version_override: Optional[str]
+    :param serialize: Boolean that indicates if identical (ie. same inputs) instances of this task should be executed in
+          serial when caching is enabled. This means that given multiple concurrent executions over identical inputs,
+          only a single instance executes and the rest wait to reuse the cached results.
+    :type serialize: bool
+    :param ignored_inputs: A tuple of input names to ignore when generating the version hash.
+    :type ignored_inputs: Union[Tuple[str, ...], str]
+    :param salt: A salt used in the hash generation.
+    :type salt: str
+    :param policies: A list of cache policies to generate the version hash.
+    :type policies: Optional[Union[List[CachePolicy], CachePolicy]]
+    """
+
+    behavior: CacheBehavior
+    version_override: Optional[str] = None
+    serialize: bool = False
+    ignored_inputs: Union[Tuple[str, ...], str] = field(default_factory=tuple)
+    salt: str = ""
+    policies: Optional[Union[List[CachePolicy], CachePolicy]] = None
+
+    def __post_init__(self):
+        if self.behavior not in get_args(CacheBehavior):
+            raise ValueError(f"Invalid cache behavior: {self.behavior}. Must be one of ['auto', 'override', 'disable']")
+
+        # Still setup _ignore_inputs when cache is disabled to prevent _ignored_inputs attribute not found error
+        if isinstance(self.ignored_inputs, str):
+            self._ignored_inputs = (self.ignored_inputs,)
+        else:
+            self._ignored_inputs = self.ignored_inputs
+
+        if self.behavior == "disable":
+            return
+
+        # Normalize policies so that self._policies is always a list
+        if self.policies is None:
+            from flyte._cache.defaults import get_default_policies
+
+            self.policies = get_default_policies()
+        elif isinstance(self.policies, CachePolicy):
+            self.policies = [self.policies]
+
+        if self.version_override is None and not self.policies:
+            raise ValueError("If version is not defined then at least one cache policy needs to be set")
+
+    def is_enabled(self) -> bool:
+        """
+        Check if the cache policy is enabled.
+        """
+        return self.behavior in ["auto", "override"]
+
+    def get_ignored_inputs(self) -> Tuple[str, ...]:
+        return self._ignored_inputs
+
+    def get_version(self, params: Optional[VersionParameters] = None) -> str:
+        if not self.is_enabled():
+            return ""
+
+        if self.version_override is not None:
+            return self.version_override
+
+        if params is None:
+            raise ValueError("Version parameters must be provided when version_override is not set.")
+
+        if params.code_bundle is not None:
+            if params.code_bundle.pkl is not None:
+                return params.code_bundle.computed_version
+
+        task_hash = ""
+        if self.policies is None:
+            raise ValueError("Cache policies are not set.")
+        policies = self.policies if isinstance(self.policies, list) else [self.policies]
+        for policy in policies:
+            try:
+                task_hash += policy.get_version(self.salt, params)
+            except Exception as e:
+                raise ValueError(f"Failed to generate version for cache policy {policy}.") from e
+
+        hash_obj = hashlib.sha256(task_hash.encode())
+        return hash_obj.hexdigest()
+
+
+CacheRequest = CacheBehavior | Cache
+
+
+def cache_from_request(cache: CacheRequest) -> Cache:
+    """
+    Coerce user input into a cache object.
+    """
+    if isinstance(cache, Cache):
+        return cache
+    return Cache(behavior=cache)

flyte/_cache/defaults.py

@@ -0,0 +1,9 @@
+from .cache import CachePolicy
+from .policy_function_body import FunctionBodyPolicy
+
+
+def get_default_policies() -> list[CachePolicy]:
+    """
+    Get default cache policies.
+    """
+    return [FunctionBodyPolicy()]

flyte/_cache/local_cache.py

@@ -0,0 +1,216 @@
+import sqlite3
+from pathlib import Path
+
+try:
+    import aiosqlite
+
+    HAS_AIOSQLITE = True
+except ImportError:
+    HAS_AIOSQLITE = False
+
+from flyteidl2.task import common_pb2
+
+from flyte._internal.runtime import convert
+from flyte._logging import logger
+from flyte.config import auto
+
+DEFAULT_CACHE_DIR = "~/.flyte"
+CACHE_LOCATION = "local-cache/cache.db"
+
+
+class LocalTaskCache(object):
+    """
+    This class implements a persistent store able to cache the result of local task executions.
+    """
+
+    _conn: "aiosqlite.Connection | None" = None
+    _conn_sync: sqlite3.Connection | None = None
+    _initialized: bool = False
+
+    @staticmethod
+    def _get_cache_path() -> str:
+        """Get the cache database path, creating directory if needed."""
+        config = auto()
+        if config.source:
+            cache_dir = config.source.parent
+        else:
+            cache_dir = Path(DEFAULT_CACHE_DIR).expanduser()
+
+        cache_path = cache_dir / CACHE_LOCATION
+        # Ensure the directory exists
+        cache_path.parent.mkdir(parents=True, exist_ok=True)
+        logger.info(f"Use local cache path: {cache_path}")
+        return str(cache_path)
+
+    @staticmethod
+    async def initialize():
+        """Initialize the cache with database connection."""
+        if not LocalTaskCache._initialized:
+            if HAS_AIOSQLITE:
+                await LocalTaskCache._initialize_async()
+            else:
+                LocalTaskCache._initialize_sync()
+
+    @staticmethod
+    async def _initialize_async():
+        """Initialize async cache connection."""
+        db_path = LocalTaskCache._get_cache_path()
+        conn = await aiosqlite.connect(db_path)
+        await conn.execute("""
+            CREATE TABLE IF NOT EXISTS task_cache (
+                key TEXT PRIMARY KEY,
+                value BLOB
+            )
+        """)
+        await conn.commit()
+        LocalTaskCache._conn = conn
+        LocalTaskCache._initialized = True
+
+    @staticmethod
+    def _initialize_sync():
+        """Initialize sync cache connection."""
+        db_path = LocalTaskCache._get_cache_path()
+        conn = sqlite3.connect(db_path)
+        conn.execute("""
+            CREATE TABLE IF NOT EXISTS task_cache (
+                key TEXT PRIMARY KEY,
+                value BLOB
+            )
+        """)
+        conn.commit()
+        LocalTaskCache._conn_sync = conn
+        LocalTaskCache._initialized = True
+
+    @staticmethod
+    async def clear():
+        """Clear all cache entries."""
+        if not LocalTaskCache._initialized:
+            await LocalTaskCache.initialize()
+
+        if HAS_AIOSQLITE:
+            await LocalTaskCache._clear_async()
+        else:
+            LocalTaskCache._clear_sync()
+
+    @staticmethod
+    async def _clear_async():
+        """Clear all cache entries (async)."""
+        if LocalTaskCache._conn is None:
+            raise RuntimeError("Cache not properly initialized")
+        await LocalTaskCache._conn.execute("DELETE FROM task_cache")
+        await LocalTaskCache._conn.commit()
+
+    @staticmethod
+    def _clear_sync():
+        """Clear all cache entries (sync)."""
+        if LocalTaskCache._conn_sync is None:
+            raise RuntimeError("Cache not properly initialized")
+        LocalTaskCache._conn_sync.execute("DELETE FROM task_cache")
+        LocalTaskCache._conn_sync.commit()
+
+    @staticmethod
+    async def get(cache_key: str) -> convert.Outputs | None:
+        if not LocalTaskCache._initialized:
+            await LocalTaskCache.initialize()
+
+        if HAS_AIOSQLITE:
+            return await LocalTaskCache._get_async(cache_key)
+        else:
+            return LocalTaskCache._get_sync(cache_key)
+
+    @staticmethod
+    async def _get_async(cache_key: str) -> convert.Outputs | None:
+        """Get cache entry (async)."""
+        if LocalTaskCache._conn is None:
+            raise RuntimeError("Cache not properly initialized")
+
+        async with LocalTaskCache._conn.execute("SELECT value FROM task_cache WHERE key = ?", (cache_key,)) as cursor:
+            row = await cursor.fetchone()
+            if row:
+                outputs_bytes = row[0]
+                outputs = common_pb2.Outputs()
+                outputs.ParseFromString(outputs_bytes)
+                return convert.Outputs(proto_outputs=outputs)
+        return None
+
+    @staticmethod
+    def _get_sync(cache_key: str) -> convert.Outputs | None:
+        """Get cache entry (sync)."""
+        if LocalTaskCache._conn_sync is None:
+            raise RuntimeError("Cache not properly initialized")
+
+        cursor = LocalTaskCache._conn_sync.execute("SELECT value FROM task_cache WHERE key = ?", (cache_key,))
+        row = cursor.fetchone()
+        if row:
+            outputs_bytes = row[0]
+            outputs = common_pb2.Outputs()
+            outputs.ParseFromString(outputs_bytes)
+            return convert.Outputs(proto_outputs=outputs)
+        return None
+
+    @staticmethod
+    async def set(
+        cache_key: str,
+        value: convert.Outputs,
+    ) -> None:
+        if not LocalTaskCache._initialized:
+            await LocalTaskCache.initialize()
+
+        if HAS_AIOSQLITE:
+            await LocalTaskCache._set_async(cache_key, value)
+        else:
+            LocalTaskCache._set_sync(cache_key, value)
+
+    @staticmethod
+    async def _set_async(
+        cache_key: str,
+        value: convert.Outputs,
+    ) -> None:
+        """Set cache entry (async)."""
+        if LocalTaskCache._conn is None:
+            raise RuntimeError("Cache not properly initialized")
+
+        output_bytes = value.proto_outputs.SerializeToString()
+        await LocalTaskCache._conn.execute(
+            "INSERT OR REPLACE INTO task_cache (key, value) VALUES (?, ?)", (cache_key, output_bytes)
+        )
+        await LocalTaskCache._conn.commit()
+
+    @staticmethod
+    def _set_sync(
+        cache_key: str,
+        value: convert.Outputs,
+    ) -> None:
+        """Set cache entry (sync)."""
+        if LocalTaskCache._conn_sync is None:
+            raise RuntimeError("Cache not properly initialized")
+
+        output_bytes = value.proto_outputs.SerializeToString()
+        LocalTaskCache._conn_sync.execute(
+            "INSERT OR REPLACE INTO task_cache (key, value) VALUES (?, ?)", (cache_key, output_bytes)
+        )
+        LocalTaskCache._conn_sync.commit()
+
+    @staticmethod
+    async def close():
+        """Close the database connection."""
+        if HAS_AIOSQLITE:
+            await LocalTaskCache._close_async()
+        else:
+            LocalTaskCache._close_sync()
+
+    @staticmethod
+    async def _close_async():
+        """Close async database connection."""
+        if LocalTaskCache._conn:
+            await LocalTaskCache._conn.close()
+            LocalTaskCache._conn = None
+        LocalTaskCache._initialized = False
+
+    @staticmethod
+    def _close_sync():
+        """Close sync database connection."""
+        if LocalTaskCache._conn_sync:
+            LocalTaskCache._conn_sync.close()
+            LocalTaskCache._conn_sync = None
+        LocalTaskCache._initialized = False

flyte/_cache/policy_function_body.py

@@ -0,0 +1,42 @@
+import ast
+import hashlib
+import inspect
+import textwrap
+
+from .cache import CachePolicy, VersionParameters
+
+
+class FunctionBodyPolicy(CachePolicy):
+    """
+    A class that implements a versioning mechanism for functions by generating
+    a SHA-256 hash of the function's source code combined with a salt.
+    """
+
+    def get_version(self, salt: str, params: VersionParameters) -> str:
+        """
+        This method generates a version string for a function by hashing the function's source code
+        combined with a salt.
+
+        :param salt: A string that is used to salt the hash.
+        :param params: VersionParameters object that contains the parameters (e.g. function, ImageSpec, etc.) that are
+                       used to generate the version.
+
+        :return: A string that represents the version of the function.
+        """
+        if params.func is None:
+            return ""
+
+        source = inspect.getsource(params.func)
+        dedented_source = textwrap.dedent(source)
+
+        # Parse the source code into an Abstract Syntax Tree (AST)
+        parsed_ast = ast.parse(dedented_source)
+
+        # Convert the AST into a string representation
+        ast_bytes = ast.dump(parsed_ast, include_attributes=False).encode("utf-8")
+
+        # Combine the AST bytes with the salt (encoded into bytes)
+        combined_data = ast_bytes + salt.encode("utf-8")
+
+        # Return the SHA-256 hash of the combined data (AST + salt)
+        return hashlib.sha256(combined_data).hexdigest()

flyte/_code_bundle/__init__.py

@@ -0,0 +1,8 @@
+from ._ignore import GitIgnore, IgnoreGroup, StandardIgnore
+from ._utils import CopyFiles
+from .bundle import build_code_bundle, build_pkl_bundle, download_bundle
+
+__all__ = ["CopyFiles", "build_code_bundle", "build_pkl_bundle", "default_ignores", "download_bundle"]
+
+
+default_ignores = [GitIgnore, StandardIgnore, IgnoreGroup]

flyte/_code_bundle/_ignore.py

@@ -0,0 +1,121 @@
+import os
+import pathlib
+import subprocess
+import tarfile as _tarfile
+from abc import ABC, abstractmethod
+from fnmatch import fnmatch
+from pathlib import Path
+from shutil import which
+from typing import List, Optional, Type
+
+from flyte._logging import logger
+
+
+class Ignore(ABC):
+    """Base for Ignores, implements core logic. Children have to implement _is_ignored"""
+
+    def __init__(self, root: Path):
+        self.root = root
+
+    def is_ignored(self, path: pathlib.Path) -> bool:
+        return self._is_ignored(path)
+
+    def tar_filter(self, tarinfo: _tarfile.TarInfo) -> Optional[_tarfile.TarInfo]:
+        if self.is_ignored(pathlib.Path(tarinfo.name)):
+            return None
+        return tarinfo
+
+    @abstractmethod
+    def _is_ignored(self, path: pathlib.Path) -> bool:
+        pass
+
+
+class GitIgnore(Ignore):
+    """Uses git cli (if available) to list all ignored files and compare with those."""
+
+    def __init__(self, root: Path):
+        super().__init__(root)
+        self.has_git = which("git") is not None
+        self.ignored_files = self._list_ignored_files()
+        self.ignored_dirs = self._list_ignored_dirs()
+
+    def _git_wrapper(self, extra_args: List[str]) -> set[str]:
+        if self.has_git:
+            out = subprocess.run(
+                ["git", "ls-files", "-io", "--exclude-standard", *extra_args],
+                cwd=self.root,
+                capture_output=True,
+                check=False,
+            )
+            if out.returncode == 0:
+                return set(out.stdout.decode("utf-8").split("\n")[:-1])
+            logger.info(f"Could not determine ignored paths due to:\n{out.stderr!r}\nNot applying any filters")
+            return set()
+        logger.info("No git executable found, not applying any filters")
+        return set()
+
+    def _list_ignored_files(self) -> set[str]:
+        return self._git_wrapper([])
+
+    def _list_ignored_dirs(self) -> set[str]:
+        return self._git_wrapper(["--directory"])
+
+    def _is_ignored(self, path: pathlib.Path) -> bool:
+        if self.ignored_files:
+            # git-ls-files uses POSIX paths
+            if Path(path).as_posix() in self.ignored_files:
+                return True
+            # Ignore empty directories
+            if os.path.isdir(os.path.join(self.root, path)) and self.ignored_dirs:
+                return Path(path).as_posix() + "/" in self.ignored_dirs
+        return False
+
+
+STANDARD_IGNORE_PATTERNS = ["*.pyc", ".cache", ".cache/*", "__pycache__", "**/__pycache__"]
+
+
+class StandardIgnore(Ignore):
+    """Retains the standard ignore functionality that previously existed. Could in theory
+    by fed with custom ignore patterns from cli."""
+
+    def __init__(self, root: Path, patterns: Optional[List[str]] = None):
+        super().__init__(root.resolve())
+        self.patterns = patterns if patterns else STANDARD_IGNORE_PATTERNS
+
+    def _is_ignored(self, path: pathlib.Path) -> bool:
+        # Convert to relative path for pattern matching
+        try:
+            rel_path = path.relative_to(self.root)
+        except ValueError:
+            # If path is not under root, don't ignore it
+            return False
+
+        for pattern in self.patterns:
+            if fnmatch(str(rel_path), pattern):
+                return True
+        return False
+
+
+class IgnoreGroup(Ignore):
+    """Groups multiple Ignores and checks a path against them. A file is ignored if any
+    Ignore considers it ignored."""
+
+    def __init__(self, root: Path, *ignores: Type[Ignore]):
+        super().__init__(root)
+        self.ignores = [ignore(root) for ignore in ignores]
+
+    def _is_ignored(self, path: pathlib.Path) -> bool:
+        for ignore in self.ignores:
+            if ignore.is_ignored(path):
+                return True
+        return False
+
+    def list_ignored(self) -> List[str]:
+        ignored = []
+        for dir, _, files in os.walk(self.root):
+            dir_path = Path(dir)
+            for file in files:
+                abs_path = dir_path / file
+                if self.is_ignored(abs_path):
+                    ignored.append(str(abs_path.relative_to(self.root)))
+        return ignored