flyte 0.0.1b0__py3-none-any.whl

union/_datastructures.py

@@ -0,0 +1,295 @@
+from __future__ import annotations
+
+import inspect
+import os
+import pathlib
+import tempfile
+from dataclasses import dataclass, field, replace
+from typing import TYPE_CHECKING, Any, Callable, Dict, Optional, Tuple, Type
+
+from union._docstring import Docstring
+from union._interface import extract_return_annotation
+from union._logging import logger
+from union._utils.helpers import base62_encode
+
+if TYPE_CHECKING:
+    from union._internal.imagebuild.image_builder import ImageCache
+    from union.report import Report
+
+
+def generate_random_name() -> str:
+    """
+    Generate a random name for the task. This is used to create unique names for tasks.
+    TODO we can use unique-namer in the future, for now its just guids
+    """
+    from uuid import uuid4
+
+    return str(uuid4())  # Placeholder for actual random name generation logic
+
+
+@dataclass(frozen=True, kw_only=True)
+class ActionID:
+    """
+    A class representing the ID of an Action, nested within a Run. This is used to identify a specific action on a task.
+    """
+
+    name: str
+    run_name: str | None = None
+    project: str | None = None
+    domain: str | None = None
+    org: str | None = None
+
+    def __post_init__(self):
+        if self.run_name is None:
+            object.__setattr__(self, "run_name", self.name)
+
+    @classmethod
+    def create_random(cls):
+        name = generate_random_name()
+        return cls(name=name, run_name=name)
+
+    def new_sub_action(self, name: str | None = None) -> ActionID:
+        """
+        Create a new sub-run with the given name. If  name is None, a random name will be generated.
+        """
+        if name is None:
+            name = generate_random_name()
+        return replace(self, name=name)
+
+    def new_sub_action_from(self, input_hash: str, group: str | None) -> ActionID:
+        """Make a deterministic name"""
+        import hashlib
+
+        components = f"{self.run_name}-{self.name}-{input_hash}" + (f"-{group}" if group else "")
+        # has the components into something deterministic
+        bytes_digest = hashlib.md5(components.encode()).digest()
+        new_name = base62_encode(bytes_digest)
+        return self.new_sub_action(new_name)
+
+
+@dataclass(frozen=True, kw_only=True)
+class RawDataPath:
+    """
+    A class representing the raw data path for a task. This is used to store the raw data for the task execution and
+    also get mutations on the path.
+    """
+
+    path: str
+
+    @classmethod
+    def from_local_folder(cls, local_folder: str | pathlib.Path | None = None) -> RawDataPath:
+        """
+        Create a new context attribute object, with local path given. Will be created if it doesn't exist.
+        :return: Path to the temporary directory
+        """
+        match local_folder:
+            case pathlib.Path():
+                local_folder.mkdir(parents=True, exist_ok=True)
+                return RawDataPath(path=str(local_folder))
+            case None:
+                # Create a temporary directory for data storage
+                p = tempfile.mkdtemp()
+                logger.debug(f"Creating temporary directory for data storage: {p}")
+                return RawDataPath(path=p)
+            case str():
+                return RawDataPath(path=local_folder)
+            case _:
+                raise ValueError(f"Invalid local path {local_folder}")
+
+    def get_random_remote_path(self, file_name: Optional[str] = None) -> str:
+        """
+        Returns a random path for uploading a file/directory to.
+
+        :param file_name: If given, will be joined after a randomly generated portion.
+        :return:
+        """
+        import random
+        from uuid import UUID
+
+        import fsspec
+        from fsspec.utils import get_protocol
+
+        random_string = UUID(int=random.getrandbits(128)).hex
+        file_prefix = self.path
+
+        protocol = get_protocol(file_prefix)
+        if "file" in protocol:
+            local_path = pathlib.Path(file_prefix) / random_string
+            if file_name:
+                # Only if file name is given do we create the parent, because it may be needed as a folder otherwise
+                local_path = local_path / file_name
+                if not local_path.exists():
+                    local_path.parent.mkdir(exist_ok=True, parents=True)
+                    local_path.touch()
+            return str(local_path.absolute())
+
+        fs = fsspec.filesystem(protocol)
+        if file_prefix.endswith(fs.sep):
+            file_prefix = file_prefix[:-1]
+        remote_path = fs.sep.join([file_prefix, random_string])
+        if file_name:
+            remote_path = fs.sep.join([remote_path, file_name])
+        return remote_path
+
+
+@dataclass(frozen=True)
+class GroupData:
+    name: str
+
+
+@dataclass(frozen=True, kw_only=True)
+class TaskContext:
+    """
+    A context class to hold the current task executions context.
+    This can be used to access various contextual parameters in the task execution by the user.
+
+    :param action: The action ID of the current execution. This is always set, within a run.
+    :param version: The version of the executed task. This is set when the task is executed by an action and will be
+      set on all sub-actions.
+    """
+
+    action: ActionID
+    version: str
+    raw_data_path: RawDataPath
+    output_path: str
+    report: Report | None = None  # This field is available if reports are enabled
+    group_data: GroupData | None = None
+    checkpoints: Checkpoints | None = None
+    code_bundle: CodeBundle | None = None
+    compiled_image_cache: ImageCache | None = None
+    data: Dict[str, Any] = field(default_factory=dict)
+
+    def replace(self, **kwargs) -> TaskContext:
+        if "data" in kwargs:
+            rec_data = kwargs.pop("data")
+            if rec_data is None:
+                return replace(self, **kwargs)
+            data = {}
+            if self.data is not None:
+                data = self.data.copy()
+            data.update(rec_data)
+            kwargs.update({"data": data})
+        return replace(self, **kwargs)
+
+    def __getitem__(self, key: str) -> Optional[Any]:
+        return self.data.get(key)
+
+
+@dataclass(frozen=True, kw_only=True)
+class CodeBundle:
+    """
+    A class representing a code bundle for a task. This is used to package the code and the inflation path.
+    The code bundle computes the version of the code using the hash of the code.
+
+    :param computed_version: The version of the code bundle. This is the hash of the code.
+    :param destination: The destination path for the code bundle to be inflated to.
+    :param tgz: Optional path to the tgz file.
+    :param pkl: Optional path to the pkl file.
+    :param downloaded_path: The path to the downloaded code bundle. This is only available during runtime, when
+        the code bundle has been downloaded and inflated.
+    """
+
+    computed_version: str
+    destination: str = "."
+    tgz: str | None = None
+    pkl: str | None = None
+    downloaded_path: pathlib.Path | None = None
+
+    # runtime_dependencies: Tuple[str, ...] = field(default_factory=tuple)  In the future if we want we could add this
+    # but this messes up actors, spark etc
+
+    def __post_init__(self):
+        if self.tgz is None and self.pkl is None:
+            raise ValueError("Either tgz or pkl must be provided")
+
+    def with_downloaded_path(self, path: pathlib.Path) -> CodeBundle:
+        """
+        Create a new CodeBundle with the given downloaded path.
+        """
+        return replace(self, downloaded_path=path)
+
+
+@dataclass(frozen=True)
+class Checkpoints:
+    """
+    A class representing the checkpoints for a task. This is used to store the checkpoints for the task execution.
+    """
+
+    prev_checkpoint_path: str | None
+    checkpoint_path: str | None
+
+
+@dataclass(frozen=True)
+class NativeInterface:
+    """
+    A class representing the native interface for a task. This is used to interact with the task and its execution
+    context.
+    """
+
+    inputs: Dict[str, Tuple[Type, Any]]
+    outputs: Dict[str, Type]
+    docstring: Optional[Docstring] = field(default=(None,))
+
+    @classmethod
+    def from_callable(cls, func: Callable) -> NativeInterface:
+        """
+        Extract the native interface from the given function. This is used to create a native interface for the task.
+        """
+        sig = inspect.signature(func)
+
+        # Extract parameter details (name, type, default value)
+        param_info = {name: (param.annotation, param.default) for name, param in sig.parameters.items()}
+
+        # Get return type
+        outputs = extract_return_annotation(sig.return_annotation)
+        return cls(inputs=param_info, outputs=outputs)
+
+    def convert_to_kwargs(self, *args, **kwargs) -> Dict[str, Any]:
+        """
+        Convert the given arguments to keyword arguments based on the native interface. This is used to convert the
+        arguments to the correct types for the task execution.
+        """
+        # Convert positional arguments to keyword arguments
+        if len(args) > len(self.inputs):
+            raise ValueError(f"Too many positional arguments provided, inputs {self.inputs.keys()}, args {len(args)}")
+        for arg, input_name in zip(args, self.inputs.keys()):
+            kwargs[input_name] = arg
+        return kwargs
+
+    def get_input_types(self) -> Dict[str, Type]:
+        """
+        Get the input types for the task. This is used to get the types of the inputs for the task execution.
+        """
+        return {k: v[0] for k, v in self.inputs.items()}
+
+
+@dataclass
+class SerializationContext:
+    """
+    This object holds serialization time contextual information, that can be used when serializing the task and
+    various parameters of a tasktemplate. This is only available when the task is being serialized and can be
+    during a deployment or runtime.
+
+    :param version: The version of the task
+    :param code_bundle: The code bundle for the task. This is used to package the code and the inflation path.
+    :param input_path: The path to the inputs for the task. This is used to determine where the inputs will be located
+    :param output_path: The path to the outputs for the task. This is used to determine where the outputs will be
+     located
+    """
+
+    version: str
+    project: str | None = None
+    domain: str | None = None
+    org: str | None = None
+    code_bundle: Optional[CodeBundle] = None
+    input_path: str = "{{.input}}"
+    output_path: str = "{{.outputPrefix}}"
+    _entrypoint_path: str = field(default="_bin/runtime.py", init=False)
+    image_cache: ImageCache | None = None
+
+    def get_entrypoint_path(self, interpreter_path: str) -> str:
+        """
+        Get the entrypoint path for the task. This is used to determine the entrypoint for the task execution.
+        :param interpreter_path: The path to the interpreter (python)
+        """
+        return os.path.join(os.path.dirname(os.path.dirname(interpreter_path)), self._entrypoint_path)

union/_deploy.py

@@ -0,0 +1,185 @@
+from __future__ import annotations
+
+import asyncio
+from dataclasses import dataclass
+from typing import TYPE_CHECKING, Dict, List, Optional, Tuple
+
+import rich.repr
+
+from ._api_commons import syncer
+from ._datastructures import SerializationContext
+from ._environment import Environment
+from ._image import Image
+from ._initialize import get_client, get_common_config, requires_client, requires_initialization
+from ._logging import logger
+from ._task import TaskTemplate
+from ._task_environment import TaskEnvironment
+
+if TYPE_CHECKING:
+    from union._protos.workflow import task_definition_pb2
+
+    from ._code_bundle import CopyFiles
+    from ._internal.imagebuild.image_builder import ImageCache
+
+
+@rich.repr.auto
+@dataclass
+class DeploymentPlan:
+    envs: Dict[str, Environment]
+    version: Optional[str] = None
+
+
+@rich.repr.auto
+@dataclass
+class Deployment:
+    envs: Dict[str, Environment]
+    deployed_tasks: List[task_definition_pb2.TaskSpec] | None = None
+
+
+@requires_client
+async def _deploy_task(
+    task: TaskTemplate, serialization_context: SerializationContext, dryrun: bool = False
+) -> task_definition_pb2.TaskSpec:
+    """
+    Deploy the given task.
+    """
+    from ._internal.runtime.task_serde import translate_task_to_wire
+    from ._protos.workflow import task_definition_pb2, task_service_pb2
+
+    image_uri = task.image.uri if isinstance(task.image, Image) else task.image
+
+    spec = translate_task_to_wire(task, serialization_context)
+    if dryrun:
+        return spec
+    logger.info(f"Deploying task {task.name}, with image {image_uri} and {serialization_context.version=} spec: {spec}")
+    cfg = get_common_config()
+    task_id = task_definition_pb2.TaskIdentifier(
+        org=cfg.org,
+        project=cfg.project,
+        domain=cfg.domain,
+        version=serialization_context.version,
+        name=spec.task_template.id.name,
+    )
+
+    resp = await get_client().task_service.DeployTask(task_service_pb2.DeployTaskRequest(task_id=task_id, spec=spec))
+    logger.info(f"Deployed task {task.name}, {resp}")
+    return spec
+
+
+async def _build_image_bg(env_name: str, image: Image) -> Tuple[str, str]:
+    """
+    Build the image in the background and return the environment name and the built image.
+    """
+    from ._build import build
+
+    logger.info(f"Building image {image.name} for environment {env_name}")
+    return env_name, await build.aio(image)
+
+
+async def build_images(deployment: DeploymentPlan) -> ImageCache:
+    """
+    Build the images for the given deployment plan and update the environment with the built image.
+    """
+    from ._internal.imagebuild.image_builder import ImageCache
+
+    images = []
+    image_identifier_map = {}
+    for env_name, env in deployment.envs.items():
+        if not isinstance(env.image, str):
+            logger.info(f"Building Image for environment {env_name}, image: {env.image}")
+            images.append(_build_image_bg(env_name, env.image))
+
+        elif env.image == "auto" and "auto" not in image_identifier_map:
+            auto_image = Image.auto()
+            image_identifier_map["auto"] = auto_image.uri
+    final_images = await asyncio.gather(*images)
+
+    for env_name, image_uri in final_images:
+        logger.info(f"Built Image for environment {env_name}, image: {image_uri}")
+        env = deployment.envs[env_name]
+        image_identifier_map[env.image.identifier] = env.image.uri
+
+    return ImageCache(image_lookup=image_identifier_map)
+
+
+@requires_initialization
+async def apply(deployment: DeploymentPlan, copy_style: CopyFiles, dryrun: bool = False) -> Deployment:
+    from ._code_bundle import build_code_bundle
+
+    cfg = get_common_config()
+    image_cache = await build_images(deployment)
+    if copy_style == "none":
+        code_bundle = None
+        assert deployment.version is not None, "Version must be set when copy_style is none"
+    else:
+        code_bundle = await build_code_bundle(from_dir=cfg.root_dir, dryrun=dryrun, copy_style=copy_style)
+        deployment.version = code_bundle.computed_version
+
+    sc = SerializationContext(
+        code_bundle=code_bundle,
+        version=deployment.version or code_bundle.computed_version,
+        image_cache=image_cache,
+    )
+
+    tasks = []
+    for env_name, env in deployment.envs.items():
+        logger.info(f"Deploying environment {env_name}")
+        if isinstance(env, TaskEnvironment):
+            for task in env.tasks.values():
+                tasks.append(_deploy_task(task, dryrun=dryrun, serialization_context=sc))
+    return Deployment(envs=deployment.envs, deployed_tasks=await asyncio.gather(*tasks))
+
+
+def _recursive_discover(
+    planned_envs: Dict[str, Environment], envs: Environment | List[Environment]
+) -> Dict[str, Environment]:
+    """
+    Recursively deploy the environment and its dependencies, if not already deployed (present in env_tasks) and
+    return the updated env_tasks.
+    """
+    if isinstance(envs, Environment):
+        envs = [envs]
+    for env in envs:
+        # Skip if the environment is already planned
+        if env.name in planned_envs:
+            continue
+        # Recursively discover dependent environments
+        for dependent_env in env.env_dep_hints:
+            _recursive_discover(planned_envs, dependent_env)
+        # Add the environment to the existing envs
+        planned_envs[env.name] = env
+    return planned_envs
+
+
+def plan_deploy(*envs: Environment, version: Optional[str] = None) -> DeploymentPlan:
+    if envs is None:
+        return DeploymentPlan({})
+    planned_envs = _recursive_discover({}, *envs)
+    return DeploymentPlan(planned_envs, version=version)
+
+
+@syncer.wrap
+async def deploy(
+    *envs: Environment,
+    dryrun: bool = False,
+    version: str | None = None,
+    interactive_mode: bool | None = None,
+    copy_style: CopyFiles = "loaded_modules",
+) -> Deployment:
+    """
+    Deploy the given environment or list of environments.
+    :param envs: Environment or list of environments to deploy.
+    :param dryrun: dryrun mode, if True, the deployment will not be applied to the control plane.
+    :param version: version of the deployment, if None, the version will be computed from the code bundle.
+    TODO: Support for interactive_mode
+    :param interactive_mode: Optional, can be forced to True or False.
+       If not provided, it will be set based on the current environment. For example Jupyter notebooks are considered
+         interactive mode, while scripts are not. This is used to determine how the code bundle is created.
+    :param copy_style: Copy style to use when running the task
+
+    :return: Deployment object containing the deployed environments and tasks.
+    """
+    if interactive_mode:
+        raise NotImplementedError("Interactive mode not yet implemented for deployment")
+    deployment = plan_deploy(*envs, version=version)
+    return await apply(deployment, copy_style=copy_style, dryrun=dryrun)

union/_doc.py

@@ -0,0 +1,29 @@
+import inspect
+from dataclasses import dataclass
+from typing import Callable
+
+
+@dataclass
+class Documentation:
+    """
+    This class is used to store the documentation of a task.
+
+    It can be set explicitly or extracted from the docstring of the task.
+    """
+
+    description: str
+
+    def __help__str__(self):
+        return self.description
+
+
+def extract_docstring(func: Callable) -> Documentation:
+    """
+    Extracts the description from a docstring.
+    """
+    if not func:
+        return Documentation(description="")
+    docstring = inspect.getdoc(func)
+    if not docstring:
+        return Documentation(description="")
+    return Documentation(description=docstring)

union/_docstring.py

@@ -0,0 +1,26 @@
+from typing import Callable, Dict, Optional
+
+
+class Docstring(object):
+    def __init__(self, docstring: Optional[str] = None, callable_: Optional[Callable] = None):
+        if docstring is not None:
+            # from docstring_parser import parse
+            self._parsed_docstring = ""  # parse(docstring)
+        else:
+            self._parsed_docstring = ""  # parse(callable_.__doc__)
+
+    @property
+    def input_descriptions(self) -> Dict[str, str]:
+        return {p.arg_name: p.description for p in self._parsed_docstring.params}
+
+    @property
+    def output_descriptions(self) -> Dict[str, str]:
+        return {p.return_name: p.description for p in self._parsed_docstring.many_returns}
+
+    @property
+    def short_description(self) -> Optional[str]:
+        return self._parsed_docstring.short_description
+
+    @property
+    def long_description(self) -> Optional[str]:
+        return self._parsed_docstring.long_description

union/_environment.py

@@ -0,0 +1,43 @@
+from __future__ import annotations
+
+from dataclasses import dataclass, field
+from typing import TYPE_CHECKING, Dict, List, Literal, Optional, Union
+
+import rich.repr
+
+from union._secret import SecretRequest
+
+from ._image import Image
+from ._resources import Resources
+
+if TYPE_CHECKING:
+    from kubernetes.client import V1PodTemplate
+
+
+@rich.repr.auto
+@dataclass(init=True, repr=True)
+class Environment:
+    """
+    :param name: Name of the environment
+    :param image: Docker image to use for the environment. If set to "auto", will use the default image.
+    :param resources: Resources to allocate for the environment.
+    :param env: Environment variables to set for the environment.
+    :param secrets: Secrets to inject into the environment.
+    :param env_dep_hints: Environment dependencies to hint, so when you deploy the environment, the dependencies are
+        also deployed. This is useful when you have a set of environments that depend on each other.
+    """
+
+    name: str
+    env_dep_hints: List[Environment] = field(default_factory=list)
+    pod_template: Optional[Union[str, "V1PodTemplate"]] = None
+    description: Optional[str] = None
+    secrets: Optional[SecretRequest] = None
+    env: Optional[Dict[str, str]] = None
+    resources: Optional[Resources] = None
+    image: Union[str, Image, Literal["auto"]] = "auto"
+
+    def add_dependency(self, *env: Environment):
+        """
+        Add a dependency to the environment.
+        """
+        self.env_dep_hints.extend(env)

union/_group.py

@@ -0,0 +1,31 @@
+from contextlib import contextmanager
+
+from ._context import internal_ctx
+from ._datastructures import GroupData
+
+
+@contextmanager
+def group(name: str):
+    """
+    Create a new group with the given name. The method is intended to be used as a context manager.
+
+    Example:
+    ```python
+    @task
+    async def my_task():
+        ...
+        with group("my_group"):
+            t1(x,y)  # tasks in this block will be grouped under "my_group"
+        ...
+    ```
+
+    :param name: The name of the group
+    """
+    ctx = internal_ctx()
+    if ctx.data.task_context is None:
+        yield
+        return
+    tctx = ctx.data.task_context
+    new_tctx = tctx.replace(group_data=GroupData(name))
+    with ctx.replace_task_context(new_tctx):
+        yield

union/_hash.py

@@ -0,0 +1,23 @@
+from typing import Callable, Generic, TypeVar
+
+T = TypeVar("T")
+
+
+class HashOnReferenceMixin(object):
+    def __hash__(self):
+        return hash(id(self))
+
+
+class HashMethod(Generic[T]):
+    """
+    Flyte-specific object used to wrap the hash function for a specific type
+    """
+
+    def __init__(self, function: Callable[[T], str]):
+        self._function = function
+
+    def calculate(self, obj: T) -> str:
+        """
+        Calculate hash for `obj`.
+        """
+        return self._function(obj)