flyte 0.0.1b0__py3-none-any.whl → 2.0.0b46__py3-none-any.whl

union/_code_bundle/bundle.py

@@ -1,176 +0,0 @@
-import asyncio
-import gzip
-import logging
-import os
-import pathlib
-import tempfile
-from pathlib import Path
-
-from flyteidl.core.tasks_pb2 import TaskTemplate
-
-import union.storage as storage
-from union._datastructures import CodeBundle
-from union._logging import log, logger
-
-from ._ignore import GitIgnore, Ignore, StandardIgnore
-from ._packaging import create_bundle, list_files_to_bundle, print_ls_tree
-from ._utils import CopyFiles, hash_file
-
-_pickled_file_extension = ".pkl.gz"
-_tar_file_extension = ".tar.gz"
-
-
-async def build_pkl_bundle(
-    o: TaskTemplate,
-    upload_to_controlplane: bool = True,
-    upload_from_dataplane_path: str | None = None,
-    copy_bundle_to: pathlib.Path | None = None,
-) -> CodeBundle:
-    """
-    Build a Pickled for the given task.
-
-    TODO We can optimize this by having an LRU cache for the function, this is so that if the same task is being
-    pickled multiple times, we can avoid the overhead of pickling it multiple times, by copying to a common place
-    and reusing based on task hash.
-
-    :param o: Object to be pickled. This is the task template.
-    :param upload_to_controlplane: Whether to upload the pickled file to the control plane or not
-    :param upload_from_dataplane_path: If we are on the dataplane, this is the path where the
-        pickled file should be uploaded to. upload_to_controlplane has to be False in this case.
-    :param copy_bundle_to: If set, the bundle will be copied to this path. This is used for testing purposes.
-    :return: CodeBundle object containing the pickled file path and the computed version.
-    """
-    import cloudpickle
-
-    import union.storage as storage
-
-    if upload_to_controlplane and upload_from_dataplane_path:
-        raise ValueError("Cannot upload to control plane and upload from dataplane path at the same time.")
-
-    logger.debug("Building pickled code bundle.")
-    with tempfile.TemporaryDirectory() as tmp_dir:
-        dest = pathlib.Path(tmp_dir) / f"code_bundle{_pickled_file_extension}"
-        with gzip.GzipFile(filename=dest, mode="wb", mtime=0) as gzipped:
-            cloudpickle.dump(o, gzipped)
-
-        if upload_to_controlplane:
-            logger.debug("Uploading pickled code bundle to control plane.")
-            from union.remote import upload_file
-
-            hash_digest, remote_path = await upload_file(dest)
-            return CodeBundle(pkl=remote_path, computed_version=hash_digest)
-
-        elif upload_from_dataplane_path:
-            logger.debug(f"Uploading pickled code bundle to dataplane path {upload_from_dataplane_path}.")
-            _, str_digest, _ = hash_file(file_path=dest)
-            final_path = await storage.put(str(dest), upload_from_dataplane_path)
-            return CodeBundle(pkl=final_path, computed_version=str_digest)
-
-        else:
-            logger.debug("Dryrun enabled, not uploading pickled code bundle.")
-            _, str_digest, _ = hash_file(file_path=dest)
-            if copy_bundle_to:
-                import shutil
-
-                # Copy the bundle to the given path
-                shutil.copy(dest, copy_bundle_to)
-                local_path = copy_bundle_to / dest.name
-                return CodeBundle(pkl=str(local_path), computed_version=str_digest)
-            return CodeBundle(pkl=str(dest), computed_version=str_digest)
-
-
-async def build_code_bundle(
-    from_dir: Path,
-    *ignore: Ignore,
-    extract_dir: str = ".",
-    dryrun: bool = False,
-    copy_bundle_to: pathlib.Path | None = None,
-    copy_style: CopyFiles = "loaded_modules",
-) -> CodeBundle:
-    """
-    Build the code bundle for the current environment.
-    :param from_dir: The directory to bundle of the code to bundle. This is the root directory for the source.
-    :param extract_dir: The directory to extract the code bundle to, when in the container. It defaults to the current
-        working directory.
-    :param ignore: The list of ignores to apply. This is a list of Ignore classes.
-    :param dryrun: If dryrun is enabled, files will not be uploaded to the control plane.
-    :param copy_bundle_to: If set, the bundle will be copied to this path. This is used for testing purposes.
-    :param copy_style: What to put into the tarball. (either all, or loaded_modules. if none, skip this function)
-
-    :return: The code bundle, which contains the path where the code was zipped to.
-    """
-    logger.debug("Building code bundle.")
-    from union.remote import upload_file
-
-    if not ignore:
-        ignore = [StandardIgnore, GitIgnore]
-
-    logger.debug(f"Finding files to bundle, ignoring as configured by: {ignore}")
-    files, digest = list_files_to_bundle(from_dir, True, *ignore, copy_style=copy_style)
-    if logger.getEffectiveLevel() <= logging.INFO:
-        print_ls_tree(from_dir, files)
-
-    logger.debug("Building code bundle.")
-    with tempfile.TemporaryDirectory() as tmp_dir:
-        bundle_path, tar_size, archive_size = create_bundle(from_dir, pathlib.Path(tmp_dir), files, digest)
-        logger.info(f"Code bundle created at {bundle_path}, size: {tar_size} MB, archive size: {archive_size} MB")
-        if not dryrun:
-            hash_digest, remote_path = await upload_file(bundle_path)
-        else:
-            remote_path = "na"
-            if copy_bundle_to:
-                import shutil
-
-                # Copy the bundle to the given path
-                shutil.copy(bundle_path, copy_bundle_to)
-                remote_path = str(copy_bundle_to / bundle_path.name)
-            _, hash_digest, _ = hash_file(file_path=bundle_path)
-        return CodeBundle(tgz=remote_path, destination=extract_dir, computed_version=hash_digest)
-
-
-@log(level=logging.INFO)
-async def download_bundle(bundle: CodeBundle) -> pathlib.Path:
-    """
-    Downloads a code bundle (tgz | pkl) to the local destination path.
-    :param bundle: The code bundle to download.
-
-    :return: The path to the downloaded code bundle.
-    """
-    dest = pathlib.Path(bundle.destination)
-    if not dest.is_dir():
-        raise ValueError(f"Destination path should be a directory, found {dest}, {dest.stat()}")
-
-    # TODO make storage apis better to accept pathlib.Path
-    if bundle.tgz:
-        downloaded_bundle = dest / os.path.basename(bundle.tgz)
-        # Download the tgz file
-        downloaded_bundle = await storage.get(bundle.tgz, str(downloaded_bundle.absolute()))
-        downloaded_bundle = pathlib.Path(downloaded_bundle)
-        # NOTE the os.path.join(destination, ''). This is to ensure that the given path is in fact a directory and all
-        # downloaded data should be copied into this directory. We do this to account for a difference in behavior in
-        # fsspec, which requires a trailing slash in case of pre-existing directory.
-        process = await asyncio.create_subprocess_exec(
-            "tar",
-            "-xvf",
-            str(downloaded_bundle),
-            "-C",
-            str(dest),
-            stdout=asyncio.subprocess.PIPE,
-            stderr=asyncio.subprocess.PIPE,
-        )
-        stdout, stderr = await process.communicate()
-
-        if process.returncode != 0:
-            raise RuntimeError(stderr.decode())
-        return downloaded_bundle.absolute()
-
-    elif bundle.pkl:
-        # Lets gunzip the pkl file
-
-        downloaded_bundle = dest / os.path.basename(bundle.pkl)
-        # Download the tgz file
-        downloaded_bundle = await storage.get(bundle.pkl, str(downloaded_bundle.absolute()))
-        downloaded_bundle = pathlib.Path(downloaded_bundle)
-        return downloaded_bundle.absolute()
-    else:
-        raise ValueError("Code bundle should be either tgz or pkl, found neither.")

union/_context.py

@@ -1,146 +0,0 @@
-from __future__ import annotations
-
-import contextvars
-from dataclasses import dataclass, replace
-from typing import TYPE_CHECKING, Callable, Optional, ParamSpec, TypeVar
-
-from union._datastructures import GroupData, RawDataPath, TaskContext
-
-if TYPE_CHECKING:
-    from union.report import Report
-
-P = ParamSpec("P")  # capture the function's parameters
-R = TypeVar("R")  # return type
-
-
-@dataclass(frozen=True, kw_only=True)
-class ContextData:
-    """
-    A ContextData cannot be created without an execution. Even for local execution's there should be an execution ID
-
-    :param: action The action ID of the current execution. This is always set, within a run.
-    :param: group_data If nested in a group the current group information
-    :param: task_context The context of the current task execution, this is what is available to the user, it is set
-        when the task is executed through `run` methods. If the Task is executed as regular python methods, this
-        will be None.
-    """
-
-    group_data: Optional[GroupData] = None
-    task_context: Optional[TaskContext] = None
-    raw_data_path: Optional[RawDataPath] = None
-
-    def replace(self, **kwargs) -> ContextData:
-        return replace(self, **kwargs)
-
-
-class Context:
-    """
-    A context class to hold the current execution context.
-    This is not coroutine safe, it assumes that the context is set in a single thread.
-    You should use the `contextual_run` function to run a function in a new context tree.
-
-    A context tree is defined as a tree of contexts, where under the root, all coroutines that were started in
-    this context tree can access the context mutations, but no coroutine, created outside of the context tree can access
-    the context mutations.
-    """
-
-    def __init__(self, data: ContextData):
-        if data is None:
-            raise ValueError("Cannot create a new context without contextdata.")
-        self._data = data
-        self._id = id(self)  # Immutable unique identifier
-
-    @property
-    def data(self) -> ContextData:
-        """Viewable data."""
-        return self._data
-
-    @property
-    def raw_data(self) -> RawDataPath:
-        """
-        Get the raw data prefix for the current context first by looking up the task context, then the raw data path
-        """
-        if self.data and self.data.task_context and self.data.task_context.raw_data_path:
-            return self.data.task_context.raw_data_path
-        if self.data and self.data.raw_data_path:
-            return self.data.raw_data_path
-        raise ValueError("Raw data path has not been set in the context.")
-
-    @property
-    def id(self) -> int:
-        """Viewable ID."""
-        return self._id
-
-    def replace_task_context(self, tctx: TaskContext) -> Context:
-        """
-        Replace the task context in the current context.
-        """
-        return Context(self.data.replace(task_context=tctx))
-
-    def new_raw_data_path(self, raw_data_path: RawDataPath) -> Context:
-        """
-        Return a copy of the context with the given raw data path object
-        """
-        return Context(self.data.replace(raw_data_path=raw_data_path))
-
-    def get_report(self) -> Optional[Report]:
-        """
-        Returns a report if within a task context, else a None
-        :return:
-        """
-        if self.data.task_context:
-            return self.data.task_context.report
-        return None
-
-    def is_task_context(self) -> bool:
-        """
-        Returns true if the context is a task context
-        :return:
-        """
-        return self.data.task_context is not None
-
-    def __enter__(self):
-        """Enter the context, setting it as the current context."""
-        self._token = root_context_var.set(self)
-        return self
-
-    def __exit__(self, exc_type, exc_val, exc_tb):
-        """Exit the context, restoring the previous context."""
-        root_context_var.reset(self._token)
-
-    async def __aenter__(self):
-        """Async version of context entry."""
-        self._token = root_context_var.set(self)
-        return self
-
-    async def __aexit__(self, exc_type, exc_val, exc_tb):
-        """Async version of context exit."""
-        root_context_var.reset(self._token)
-
-    def __repr__(self):
-        return f"{self.data}"
-
-    def __str__(self):
-        return self.__repr__()
-
-
-# Global context variable to hold the current context
-root_context_var = contextvars.ContextVar("root", default=Context(data=ContextData()))
-
-
-def ctx() -> Optional[TaskContext]:
-    """Retrieve the current task context from the context variable."""
-    return internal_ctx().data.task_context
-
-
-def internal_ctx() -> Context:
-    """Retrieve the current context from the context variable."""
-    return root_context_var.get()
-
-
-async def contextual_run(func: Callable[P, R], *args: P.args, **kwargs: P.kwargs) -> R:
-    """
-    Run a function with a new context subtree.
-    """
-    _ctx = contextvars.copy_context()
-    return await _ctx.run(func, *args, **kwargs)

union/_datastructures.py

@@ -1,295 +0,0 @@
-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)