flyte 0.0.1b0__py3-none-any.whl

flyte/__init__.py

@@ -0,0 +1,62 @@
+"""
+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__ = [
+    "ABFS",
+    "GCS",
+    "GPU",
+    "S3",
+    "TPU",
+    "Cache",
+    "CachePolicy",
+    "CacheRequest",
+    "Device",
+    "Image",
+    "Resources",
+    "RetryStrategy",
+    "ReusePolicy",
+    "Secret",
+    "SecretRequest",
+    "TaskEnvironment",
+    "Timeout",
+    "TimeoutType",
+    "__version__",
+    "ctx",
+    "deploy",
+    "group",
+    "init",
+    "run",
+    "trace",
+    "with_runcontext",
+]
+
+from ._cache import Cache, CachePolicy, CacheRequest
+from ._context import ctx
+from ._deploy import deploy
+from ._group import group
+from ._image import Image
+from ._initialize import ABFS, GCS, S3, init
+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__

flyte/_api_commons.py

@@ -0,0 +1,3 @@
+from synchronicity import Synchronizer
+
+syncer = Synchronizer()

flyte/_bin/runtime.py

@@ -0,0 +1,126 @@
+"""
+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 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"
+
+
+@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._utils as utils
+    from flyte._datastructures import ActionID, Checkpoints, CodeBundle, RawDataPath
+    from flyte._initialize import S3, initialize_in_cluster
+    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
+
+    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", "")
+
+    ep = os.environ.get(ENDPOINT_OVERRIDE, "host.docker.internal:8090")
+
+    bundle = CodeBundle(tgz=tgz, pkl=pkl, destination=dest, computed_version=version)
+    # TODO configure storage correctly for cluster
+    initialize_in_cluster(storage=S3.auto())
+    controller = create_controller(ct="remote", endpoint=ep, insecure=True)
+
+    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,25 @@
+from __future__ import annotations
+
+from ._api_commons import syncer
+from ._image import Image
+
+
+@syncer.wrap
+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,146 @@
+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
+
+from flyte._datastructures import CodeBundle
+
+# if TYPE_CHECKING:
+from flyte._image import Image
+
+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()