flyte 0.0.1b0__py3-none-any.whl

union/_cache/cache.py

@@ -0,0 +1,141 @@
+import hashlib
+from dataclasses import dataclass, field
+from typing import (
+    Callable,
+    Generic,
+    List,
+    Optional,
+    Protocol,
+    Tuple,
+    Union,
+    runtime_checkable,
+)
+
+from typing_extensions import Literal, ParamSpec, TypeVar
+
+from union._datastructures import CodeBundle
+
+# if TYPE_CHECKING:
+from union._image import Image
+
+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]
+    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: ...
+
+
+@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 ["auto", "override", "disable"]:
+            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 union._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 = ""
+        for policy in self.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)

union/_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()]

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

union/_cli/_common.py

@@ -0,0 +1,263 @@
+from __future__ import annotations
+
+import importlib.util
+import logging
+import os
+import sys
+from abc import abstractmethod
+from dataclasses import Field, dataclass, field, replace
+from pathlib import Path
+from types import MappingProxyType, ModuleType
+from typing import Any, Dict, Iterable, List, Optional
+
+import rich_click as click
+from rich.table import Table
+
+import union.errors
+
+PROJECT_OPTION = click.Option(
+    param_decls=["-p", "--project"],
+    required=False,
+    type=str,
+    default="default",
+    help="Project to operate on",
+    show_default=True,
+)
+
+DOMAIN_OPTION = click.Option(
+    param_decls=["-d", "--domain"],
+    required=False,
+    type=str,
+    default="development",
+    help="Domain to operate on",
+    show_default=True,
+)
+
+DRY_RUN_OPTION = click.Option(
+    param_decls=["--dry-run", "--dryrun"],
+    required=False,
+    type=bool,
+    is_flag=True,
+    default=False,
+    help="Dry run, do not actually call the backend service.",
+    show_default=True,
+)
+
+
+def _common_options() -> List[click.Option]:
+    """
+    Common options for that will be added to all commands and groups that inherit from CommandBase or GroupBase.
+    """
+    return [PROJECT_OPTION, DOMAIN_OPTION]
+
+
+# This is global state for the CLI, it is manipulated by the main command
+
+
+@dataclass(frozen=True)
+class CLIConfig:
+    """
+    This is the global state for the CLI. It is manipulated by the main command.
+    """
+
+    log_level: int = logging.ERROR
+    endpoint: str = "dns:///dogfood.cloud-staging.union.ai"
+    insecure: bool = False
+    org_override: str | None = None
+
+    def replace(self, **kwargs) -> CLIConfig:
+        """
+        Replace the global state with a new one.
+        """
+        return replace(self, **kwargs)
+
+    def init(self, project: str | None = None, domain: str | None = None):
+        import union
+
+        union.init(
+            endpoint=self.endpoint,
+            project=project,
+            domain=domain,
+            log_level=self.log_level,
+            insecure=self.insecure,
+            org=self.org_override,
+        )
+
+
+class InvokeBaseMixin:
+    """
+    Mixin to catch grpc.RpcError, union.RpcError, other errors and other exceptions and raise them as
+     gclick.ClickException.
+    """
+
+    def invoke(self, ctx):
+        import grpc
+
+        try:
+            return super().invoke(ctx)  # type: ignore
+        except grpc.aio.AioRpcError as e:
+            if e.code() == grpc.StatusCode.UNAUTHENTICATED:
+                raise click.ClickException(f"Authentication failed. Please check your credentials. {e!s}")
+            if e.code() == grpc.StatusCode.NOT_FOUND:
+                raise click.ClickException(f"Requested object NOT FOUND. Please check your input. Error: {e!s}")
+            if e.code() == grpc.StatusCode.ALREADY_EXISTS:
+                raise click.ClickException("Resource already exists.")
+            raise click.ClickException(f"Error invoking command: {e!s}") from e
+        except union.errors.InitializationError:
+            raise click.ClickException("Initialize the CLI with a remote configuration. For example, pass --endpoint")
+        except Exception as e:
+            raise click.ClickException(f"Error invoking command: {e}") from e
+
+
+class CommandBase(InvokeBaseMixin, click.RichCommand):
+    """
+    Base class for all commands, that adds common options to all commands if enabled.
+    """
+
+    common_options_enabled = True
+
+    def __init__(self, *args, **kwargs):
+        if "params" not in kwargs:
+            kwargs["params"] = []
+        if self.common_options_enabled:
+            kwargs["params"].extend(_common_options())
+        super().__init__(*args, **kwargs)
+
+
+class GroupBase(InvokeBaseMixin, click.RichGroup):
+    """
+    Base class for all commands, that adds common options to all commands if enabled.
+    """
+
+    common_options_enabled = True
+
+    def __init__(self, *args, **kwargs):
+        if "params" not in kwargs:
+            kwargs["params"] = []
+        if self.common_options_enabled:
+            kwargs["params"].extend(_common_options())
+        super().__init__(*args, **kwargs)
+
+
+class GroupBaseNoOptions(GroupBase):
+    common_options_enabled = False
+
+
+def make_click_option_field(o: click.Option) -> Field:
+    if o.multiple:
+        o.help = click.style("Multiple values allowed.", bold=True) + f"{o.help}"
+        return field(default_factory=lambda: o.default, metadata={"click.option": o})
+    return field(default=o.default, metadata={"click.option": o})
+
+
+def get_option_from_metadata(metadata: MappingProxyType) -> click.Option:
+    return metadata["click.option"]
+
+
+def key_value_callback(_: Any, param: str, values: List[str]) -> Optional[Dict[str, str]]:
+    """
+    Callback for click to parse key-value pairs.
+    """
+    if not values:
+        return None
+    result = {}
+    for v in values:
+        if "=" not in v:
+            raise click.BadParameter(f"Expected key-value pair of the form key=value, got {v}")
+        k, v_ = v.split("=", 1)
+        result[k.strip()] = v_.strip()
+    return result
+
+
+class ObjectsPerFileGroup(GroupBase):
+    """
+    Group that creates a command for each object in a python file.
+    """
+
+    def __init__(self, filename: Path, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+        if not filename.exists():
+            raise click.ClickException(f"{filename} does not exists")
+        self.filename = filename
+        self._objs = None
+
+    @abstractmethod
+    def _filter_objects(self, module: ModuleType) -> Dict[str, Any]:
+        raise NotImplementedError
+
+    @property
+    def objs(self) -> Dict[str, Any]:
+        if self._objs is not None:
+            return self._objs
+
+        module_name = os.path.splitext(os.path.basename(self.filename))[0]
+        module_path = os.path.dirname(os.path.abspath(self.filename))
+
+        spec = importlib.util.spec_from_file_location(module_name, self.filename)
+        module = importlib.util.module_from_spec(spec)
+        sys.modules[module_name] = module
+
+        sys.path.append(module_path)
+        spec.loader.exec_module(module)
+
+        self._objs = self._filter_objects(module)
+        return self._objs
+
+    def list_commands(self, ctx):
+        m = list(self.objs.keys())
+        return sorted(m)
+
+    @abstractmethod
+    def _get_command_for_obj(self, ctx: click.Context, obj_name: str, obj: Any) -> click.Command: ...
+
+    def get_command(self, ctx, obj_name):
+        obj = self.objs[obj_name]
+        return self._get_command_for_obj(ctx, obj_name, obj)
+
+
+class FileGroup(GroupBase):
+    """
+    Group that creates a command for each file in the current directory that is not __init__.py.
+    """
+
+    common_options_enabled = False
+
+    def __init__(
+        self,
+        *args,
+        directory: Path | None = None,
+        **kwargs,
+    ):
+        if "params" not in kwargs:
+            kwargs["params"] = []
+        super().__init__(*args, **kwargs)
+        self._files = None
+        self._dir = directory
+
+    @property
+    def files(self):
+        if self._files is None:
+            directory = self._dir or Path(".").absolute()
+            self._files = [os.fspath(p) for p in directory.glob("*.py") if p.name != "__init__.py"]
+        return self._files
+
+    def list_commands(self, ctx):
+        return self.files
+
+    def get_command(self, ctx, filename):
+        raise NotImplementedError
+
+
+def get_table(title: str, vals: Iterable[Any]) -> Table:
+    """
+    Get a table from a list of values.
+    """
+    table = Table(title=title)
+    headers = None
+    for p in vals:
+        if headers is None:
+            headers = [k for k, _ in p.__rich_repr__()]
+            for h in headers:
+                table.add_column(h.capitalize())
+        table.add_row(*[str(v) for _, v in p.__rich_repr__()])
+    return table

union/_cli/_create.py

@@ -0,0 +1,40 @@
+from typing import get_args
+
+import rich_click as click
+
+import union._cli._common as common
+from union.remote._secret import SecretTypes
+
+
+@click.group(name="create")
+def create():
+    """
+    Create a new task or environment.
+    """
+
+
+@create.command(cls=common.CommandBase)
+@click.argument("name", type=str, required=True)
+@click.argument("value", type=str, required=False)
+@click.option("--from-file", type=click.Path(exists=True), help="Path to the file with the binary secret.")
+@click.option("--type", type=click.Choice(get_args(SecretTypes)), default="regular", help="Type of the secret.")
+@click.pass_obj
+def secret(
+    cfg: common.CLIConfig,
+    name: str,
+    value: str | None = None,
+    from_file: str | None = None,
+    type: SecretTypes = "regular",
+    project: str | None = None,
+    domain: str | None = None,
+):
+    """
+    Create a new secret.
+    """
+    from union.remote import Secret
+
+    cfg.init(project, domain)
+    if from_file:
+        with open(from_file, "rb") as f:
+            value = f.read()
+    Secret.create(name=name, value=value, type=type)

union/_cli/_delete.py

@@ -0,0 +1,23 @@
+import rich_click as click
+
+import union._cli._common as common
+
+
+@click.group(name="delete")
+def delete():
+    """
+    Delete a task or environment.
+    """
+
+
+@click.command(cls=common.CommandBase)
+@click.argument("name", type=str, required=True)
+@click.pass_obj
+def secret(cfg: common.CLIConfig, name: str, project: str | None = None, domain: str | None = None):
+    """
+    Delete a secret.
+    """
+    from union.remote import Secret
+
+    cfg.init(project, domain)
+    Secret.delete(name=name)

union/_cli/_deploy.py

@@ -0,0 +1,120 @@
+from dataclasses import dataclass, fields
+from pathlib import Path
+from types import ModuleType
+from typing import Any, Dict, List, cast
+
+import click
+
+import union
+
+from .._code_bundle._utils import CopyFiles
+from . import _common as common
+from ._common import CLIConfig
+
+
+@dataclass
+class DeployArguments:
+    project: str = common.make_click_option_field(common.PROJECT_OPTION)
+    domain: str = common.make_click_option_field(common.DOMAIN_OPTION)
+    dry_run: bool = common.make_click_option_field(common.DRY_RUN_OPTION)
+    copy_style: CopyFiles = common.make_click_option_field(
+        click.Option(
+            ["--copy-style"],
+            type=click.Choice(["loaded_modules", "all_files", "none"]),
+            default="loaded_modules",
+            help="Copy style to use when running the task",
+        )
+    )
+    version: CopyFiles = common.make_click_option_field(
+        click.Option(
+            ["--version"],
+            type=str,
+            required=False,
+            help="If no files are copied, user must specify a version to register tasks with",
+        )
+    )
+
+    @classmethod
+    def from_dict(cls, d: Dict[str, Any]) -> "DeployArguments":
+        return cls(**d)
+
+    @classmethod
+    def options(cls) -> List[click.Option]:
+        """
+        Return the set of base parameters added to every pyflyte run workflow subcommand.
+        """
+        return [common.get_option_from_metadata(f.metadata) for f in fields(cls) if f.metadata]
+
+
+class DeployEnvCommand(click.Command):
+    def __init__(self, obj_name: str, obj: Any, deploy_args: DeployArguments, *args, **kwargs):
+        self.obj_name = obj_name
+        self.obj = obj
+        self.deploy_args = deploy_args
+        super().__init__(name=obj_name, *args, **kwargs)
+
+    def invoke(self, ctx):
+        print(f"Deploying environment: {self.obj_name}")
+        obj: CLIConfig = ctx.obj
+        obj.remote_init(self.deploy_args.project, self.deploy_args.domain)
+        return union.deploy(
+            self.obj,
+            dryrun=self.deploy_args.dry_run,
+            copy_style=self.deploy_args.copy_style,
+            version=self.deploy_args.version,
+        )
+
+
+class EnvPerFileGroup(common.ObjectsPerFileGroup):
+    """
+    Group that creates a command for each task in the current directory that is not __init__.py.
+    """
+
+    def __init__(self, filename: Path, deploy_args: DeployArguments, *args, **kwargs):
+        super().__init__(*args, filename=filename, **kwargs)
+        self.deploy_args = deploy_args
+
+    def _filter_objects(self, module: ModuleType) -> Dict[str, Any]:
+        return {k: v for k, v in module.__dict__.items() if isinstance(v, union.TaskEnvironment)}
+
+    def _get_command_for_obj(self, ctx: click.Context, obj_name: str, obj: Any) -> click.Command:
+        obj = cast(union.TaskEnvironment, obj)
+        return DeployEnvCommand(
+            obj_name=obj_name,
+            obj=obj,
+            help=obj.description,
+            deploy_args=self.deploy_args,
+        )
+
+
+class EnvFiles(common.FileGroup):
+    """
+    Group that creates a command for each file in the current directory that is not __init__.py.
+    """
+
+    common_options_enabled = False
+
+    def __init__(
+        self,
+        *args,
+        **kwargs,
+    ):
+        if "params" not in kwargs:
+            kwargs["params"] = []
+        kwargs["params"].extend(DeployArguments.options())
+        super().__init__(*args, **kwargs)
+
+    def get_command(self, ctx, filename):
+        deploy_args = DeployArguments.from_dict(ctx.params)
+        return EnvPerFileGroup(
+            filename=Path(filename),
+            deploy_args=deploy_args,
+            name=filename,
+            help=f"Run, functions decorated `env.task` or instances of Tasks in {filename}",
+        )
+
+
+deploy = EnvFiles(
+    name="deploy",
+    help="deploy one or more environments from a python file.",
+)