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

flyte/__init__.py

@@ -1,2 +1,78 @@
-def hello() -> str:
-    return "Hello from flyte!"
+"""
+Flyte SDK for authoring Compound AI applications, services and workflows.
+
+## Environments
+
+TaskEnvironment class to define a new environment for a set of tasks.
+
+Example usage:
+
+```python
+env = flyte.TaskEnvironment(name="my_env", image="my_image", resources=Resources(cpu="1", memory="1Gi"))
+
+@env.task
+async def my_task():
+    pass
+```
+"""
+
+__all__ = [
+    "GPU",
+    "TPU",
+    "Cache",
+    "CachePolicy",
+    "CacheRequest",
+    "Device",
+    "Environment",
+    "Image",
+    "PodTemplate",
+    "Resources",
+    "RetryStrategy",
+    "ReusePolicy",
+    "Secret",
+    "SecretRequest",
+    "TaskEnvironment",
+    "Timeout",
+    "TimeoutType",
+    "__version__",
+    "ctx",
+    "deploy",
+    "group",
+    "init",
+    "init_from_config",
+    "map",
+    "run",
+    "trace",
+    "with_runcontext",
+]
+
+import sys
+
+from ._cache import Cache, CachePolicy, CacheRequest
+from ._context import ctx
+from ._deploy import deploy
+from ._environment import Environment
+from ._excepthook import custom_excepthook
+from ._group import group
+from ._image import Image
+from ._initialize import init, init_from_config
+from ._map import map
+from ._pod import PodTemplate
+from ._resources import GPU, TPU, Device, Resources
+from ._retry import RetryStrategy
+from ._reusable_environment import ReusePolicy
+from ._run import run, with_runcontext
+from ._secret import Secret, SecretRequest
+from ._task_environment import TaskEnvironment
+from ._timeout import Timeout, TimeoutType
+from ._trace import trace
+from ._version import __version__
+
+sys.excepthook = custom_excepthook
+
+
+def version() -> str:
+    """
+    Returns the version of the Flyte SDK.
+    """
+    return __version__

flyte/_bin/runtime.py

@@ -0,0 +1,152 @@
+"""
+Flyte runtime module, this is the entrypoint script for the Flyte runtime.
+
+Caution: Startup time for this module is very important, as it is the entrypoint for the Flyte runtime.
+Refrain from importing any modules here. If you need to import any modules, do it inside the main function.
+"""
+
+import asyncio
+import os
+import sys
+from typing import Any, List
+
+import click
+
+# Todo: work with pvditt to make these the names
+# ACTION_NAME = "_U_ACTION_NAME"
+# RUN_NAME = "_U_RUN_NAME"
+# PROJECT_NAME = "_U_PROJECT_NAME"
+# DOMAIN_NAME = "_U_DOMAIN_NAME"
+# ORG_NAME = "_U_ORG_NAME"
+
+ACTION_NAME = "ACTION_NAME"
+RUN_NAME = "RUN_NAME"
+PROJECT_NAME = "FLYTE_INTERNAL_TASK_PROJECT"
+DOMAIN_NAME = "FLYTE_INTERNAL_TASK_DOMAIN"
+ORG_NAME = "_U_ORG_NAME"
+ENDPOINT_OVERRIDE = "_U_EP_OVERRIDE"
+RUN_OUTPUT_BASE_DIR = "_U_RUN_BASE"
+
+# TODO: Remove this after proper auth is implemented
+_UNION_EAGER_API_KEY_ENV_VAR = "_UNION_EAGER_API_KEY"
+
+
+@click.command("a0")
+@click.option("--inputs", "-i", required=True)
+@click.option("--outputs-path", "-o", required=True)
+@click.option("--version", "-v", required=True)
+@click.option("--run-base-dir", envvar=RUN_OUTPUT_BASE_DIR, required=True)
+@click.option("--raw-data-path", "-r", required=False)
+@click.option("--checkpoint-path", "-c", required=False)
+@click.option("--prev-checkpoint", "-p", required=False)
+@click.option("--name", envvar=ACTION_NAME, required=False)
+@click.option("--run-name", envvar=RUN_NAME, required=False)
+@click.option("--project", envvar=PROJECT_NAME, required=False)
+@click.option("--domain", envvar=DOMAIN_NAME, required=False)
+@click.option("--org", envvar=ORG_NAME, required=False)
+@click.option("--image-cache", required=False)
+@click.option("--tgz", required=False)
+@click.option("--pkl", required=False)
+@click.option("--dest", required=False)
+@click.option("--resolver", required=False)
+@click.argument(
+    "resolver-args",
+    type=click.UNPROCESSED,
+    nargs=-1,
+)
+def main(
+    run_name: str,
+    name: str,
+    project: str,
+    domain: str,
+    org: str,
+    image_cache: str,
+    version: str,
+    inputs: str,
+    run_base_dir: str,
+    outputs_path: str,
+    raw_data_path: str,
+    checkpoint_path: str,
+    prev_checkpoint: str,
+    tgz: str,
+    pkl: str,
+    dest: str,
+    resolver: str,
+    resolver_args: List[str],
+):
+    sys.path.insert(0, ".")
+
+    import flyte
+    import flyte._utils as utils
+    from flyte._initialize import init
+    from flyte._internal.controllers import create_controller
+    from flyte._internal.imagebuild.image_builder import ImageCache
+    from flyte._internal.runtime.entrypoints import load_and_run_task
+    from flyte._logging import logger
+    from flyte.models import ActionID, Checkpoints, CodeBundle, RawDataPath
+
+    logger.info(f"Initializing flyte runtime - version {flyte.__version__}")
+
+    assert org, "Org is required for now"
+    assert project, "Project is required"
+    assert domain, "Domain is required"
+    assert run_name, f"Run name is required {run_name}"
+    assert name, f"Action name is required {name}"
+
+    if run_name.startswith("{{"):
+        run_name = os.getenv("RUN_NAME", "")
+    if name.startswith("{{"):
+        name = os.getenv("ACTION_NAME", "")
+
+    # Figure out how to connect
+    # This detection of api key is a hack for now.
+    controller_kwargs: dict[str, Any] = {"insecure": False}
+    if api_key := os.getenv(_UNION_EAGER_API_KEY_ENV_VAR):
+        logger.info("Using api key from environment")
+        controller_kwargs["api_key"] = api_key
+    else:
+        ep = os.environ.get(ENDPOINT_OVERRIDE, "host.docker.internal:8090")
+        controller_kwargs["endpoint"] = ep
+        if "localhost" in ep or "docker" in ep:
+            controller_kwargs["insecure"] = True
+        logger.debug(f"Using controller endpoint: {ep} with kwargs: {controller_kwargs}")
+
+    bundle = CodeBundle(tgz=tgz, pkl=pkl, destination=dest, computed_version=version)
+    # We init regular client here so that reference tasks can work
+    # Current reference tasks will not work with remote controller, because we create 2 different
+    # channels on different threads and this is not supported by grpcio or the auth system. It ends up leading
+    # File "src/python/grpcio/grpc/_cython/_cygrpc/aio/completion_queue.pyx.pxi", line 147,
+    # in grpc._cython.cygrpc.PollerCompletionQueue._handle_events
+    # BlockingIOError: [Errno 11] Resource temporarily unavailable
+    # init(org=org, project=project, domain=domain, **controller_kwargs)
+    # TODO solution is to use a single channel for both controller and reference tasks, but this requires a refactor
+    init()
+    # Controller is created with the same kwargs as init, so that it can be used to run tasks
+    controller = create_controller(ct="remote", **controller_kwargs)
+
+    ic = ImageCache.from_transport(image_cache) if image_cache else None
+
+    # Create a coroutine to load the task and run it
+    task_coroutine = load_and_run_task(
+        resolver=resolver,
+        resolver_args=resolver_args,
+        action=ActionID(name=name, run_name=run_name, project=project, domain=domain, org=org),
+        raw_data_path=RawDataPath(path=raw_data_path),
+        checkpoints=Checkpoints(checkpoint_path, prev_checkpoint),
+        code_bundle=bundle,
+        input_path=inputs,
+        output_path=outputs_path,
+        run_base_dir=run_base_dir,
+        version=version,
+        controller=controller,
+        image_cache=ic,
+    )
+    # Create a coroutine to watch for errors
+    controller_failure = controller.watch_for_errors()
+
+    # Run both coroutines concurrently and wait for first to finish and cancel the other
+    async def _run_and_stop():
+        await utils.run_coros(controller_failure, task_coroutine)
+        await controller.stop()
+
+    asyncio.run(_run_and_stop())

flyte/_build.py

@@ -0,0 +1,26 @@
+from __future__ import annotations
+
+from flyte.syncify import syncify
+
+from ._image import Image
+
+
+@syncify
+async def build(image: Image) -> str:
+    """
+    Build an image. The existing async context will be used.
+
+    Example:
+    ```
+    import flyte
+    image = flyte.Image("example_image")
+    if __name__ == "__main__":
+        asyncio.run(flyte.build.aio(image))
+    ```
+
+    :param image: The image(s) to build.
+    :return: The image URI.
+    """
+    from flyte._internal.imagebuild.image_builder import ImageBuildEngine
+
+    return await ImageBuildEngine.build(image)

flyte/_cache/__init__.py

@@ -0,0 +1,12 @@
+from .cache import Cache, CacheBehavior, CachePolicy, CacheRequest
+from .defaults import get_default_policies
+from .policy_function_body import FunctionBodyPolicy
+
+__all__ = [
+    "Cache",
+    "CacheBehavior",
+    "CachePolicy",
+    "CacheRequest",
+    "FunctionBodyPolicy",
+    "get_default_policies",
+]

flyte/_cache/cache.py

@@ -0,0 +1,145 @@
+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", "enabled"]
+
+
+@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']")
+        if self.behavior == "disable":
+            return
+
+        if isinstance(self.ignored_inputs, str):
+            self._ignored_inputs = (self.ignored_inputs,)
+        else:
+            self._ignored_inputs = self.ignored_inputs
+
+        # 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/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,113 @@
+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)
+        self.patterns = patterns if patterns else STANDARD_IGNORE_PATTERNS
+
+    def _is_ignored(self, path: pathlib.Path) -> bool:
+        for pattern in self.patterns:
+            if fnmatch(str(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 self.root.walk():
+            for file in files:
+                abs_path = dir / file
+                if self.is_ignored(abs_path):
+                    ignored.append(str(abs_path.relative_to(self.root)))
+        return ignored