flyte 0.0.1b0__py3-none-any.whl

union/_code_bundle/_utils.py

@@ -0,0 +1,342 @@
+from __future__ import annotations
+
+import gzip
+import hashlib
+import importlib.util
+import os
+import pathlib
+import shutil
+import site
+import stat
+import sys
+import tarfile
+import tempfile
+import typing
+from datetime import datetime, timezone
+from functools import lru_cache
+from pathlib import Path
+from types import ModuleType
+from typing import List, Literal, Optional, Tuple, Union
+
+from union._logging import logger
+
+from ._ignore import IgnoreGroup
+
+CopyFiles = Literal["loaded_modules", "all", "none"]
+
+
+def compress_scripts(source_path: str, destination: str, modules: List[ModuleType]):
+    """
+    Compresses the single script while maintaining the folder structure for that file.
+
+    For example, given the follow file structure:
+    .
+    ├── flyte
+              ├── __init__.py
+              └── workflows
+                  ├── example.py
+                  ├── another_example.py
+                  ├── yet_another_example.py
+                  ├── unused_example.py
+                  └── __init__.py
+
+    Let's say you want to compress `example.py` imports `another_example.py`. And `another_example.py`
+    imports on `yet_another_example.py`. This will  produce a tar file that contains only that
+    file alongside with the folder structure, i.e.:
+
+    .
+    ├── flyte
+              ├── __init__.py
+              └── workflows
+                  ├── example.py
+                  ├── another_example.py
+                  ├── yet_another_example.py
+                  └── __init__.py
+
+    """
+    with tempfile.TemporaryDirectory() as tmp_dir:
+        destination_path = os.path.join(tmp_dir, "code")
+        os.mkdir(destination_path)
+        add_imported_modules_from_source(source_path, destination_path, modules)
+
+        tar_path = os.path.join(tmp_dir, "tmp.tar")
+        with tarfile.open(tar_path, "w") as tar:
+            tmp_path: str = os.path.join(tmp_dir, "code")
+            files: typing.List[str] = os.listdir(tmp_path)
+            for ws_file in files:
+                tar.add(os.path.join(tmp_path, ws_file), arcname=ws_file, filter=tar_strip_file_attributes)
+        with gzip.GzipFile(filename=destination, mode="wb", mtime=0) as gzipped:
+            with open(tar_path, "rb") as tar_file:
+                gzipped.write(tar_file.read())
+
+
+# Takes in a TarInfo and returns the modified TarInfo:
+# https://docs.python.org/3/library/tarfile.html#tarinfo-objects
+# intended to be passed as a filter to tarfile.add
+# https://docs.python.org/3/library/tarfile.html#tarfile.TarFile.add
+def tar_strip_file_attributes(tar_info: tarfile.TarInfo) -> tarfile.TarInfo:
+    # set time to epoch timestamp 0, aka 00:00:00 UTC on 1 January 1980
+    # note that when extracting this tarfile, this time will be shown as the modified date
+    tar_info.mtime = datetime(1980, 1, 1, tzinfo=timezone.utc).timestamp()
+
+    # user/group info
+    tar_info.uid = 0
+    tar_info.uname = ""
+    tar_info.gid = 0
+    tar_info.gname = ""
+
+    # stripping paxheaders may not be required
+    # see https://stackoverflow.com/questions/34688392/paxheaders-in-tarball
+    tar_info.pax_headers = {}
+
+    return tar_info
+
+
+def ls_files(
+    source_path: pathlib.Path,
+    copy_file_detection: CopyFiles,
+    deref_symlinks: bool = False,
+    ignore_group: Optional[IgnoreGroup] = None,
+) -> Tuple[List[str], str]:
+    """
+    user_modules_and_packages is a list of the Python modules and packages, expressed as absolute paths, that the
+    user has run this pyflyte command with. For pyflyte run for instance, this is just a list of one.
+    This is used for two reasons.
+      - Everything in this list needs to be returned. Files are returned and folders are walked.
+      - A common source path is derived from this is, which is just the common folder that contains everything in the
+        list. For ex. if you do
+        $ pyflyte --pkgs a.b,a.c package
+        Then the common root is just the folder a/. The modules list is filtered against this root. Only files
+        representing modules under this root are included
+
+    If the copy enum is set to loaded_modules, then the loaded sys modules will be used.
+    """
+
+    # Unlike the below, the value error here is useful and should be returned to the user, like if absolute and
+    # relative paths are mixed.
+
+    # This is --copy auto
+    if copy_file_detection == "loaded_modules":
+        sys_modules = list(sys.modules.values())
+        all_files = list_imported_modules_as_files(str(source_path), sys_modules)
+    # this is --copy all (--copy none should never invoke this function)
+    else:
+        all_files = list_all_files(source_path, deref_symlinks, ignore_group)
+
+    all_files.sort()
+    hasher = hashlib.md5()
+    for abspath in all_files:
+        relpath = os.path.relpath(abspath, source_path)
+        _filehash_update(abspath, hasher)
+        _pathhash_update(relpath, hasher)
+
+    digest = hasher.hexdigest()
+
+    return all_files, digest
+
+
+def _filehash_update(path: Union[os.PathLike, str], hasher: hashlib._Hash) -> None:
+    blocksize = 65536
+    with open(path, "rb") as f:
+        bytes = f.read(blocksize)
+        while bytes:
+            hasher.update(bytes)
+            bytes = f.read(blocksize)
+
+
+def _pathhash_update(path: Union[os.PathLike, str], hasher: hashlib._Hash) -> None:
+    path_list = path.split(os.sep)
+    hasher.update("".join(path_list).encode("utf-8"))
+
+
+EXCLUDE_DIRS = {".git"}
+
+
+def list_all_files(source_path: pathlib.Path, deref_symlinks, ignore_group: Optional[IgnoreGroup] = None) -> List[str]:
+    all_files = []
+
+    # This is needed to prevent infinite recursion when walking with followlinks
+    visited_inodes = set()
+    for root, dirnames, files in source_path.walk(top_down=True, follow_symlinks=deref_symlinks):
+        dirnames[:] = [d for d in dirnames if d not in EXCLUDE_DIRS]
+        if deref_symlinks:
+            inode = os.stat(root).st_ino
+            if inode in visited_inodes:
+                continue
+            visited_inodes.add(inode)
+
+        ff = []
+        files.sort()
+        for fname in files:
+            abspath = (root / fname).absolute()
+            # Only consider files that exist (e.g. disregard symlinks that point to non-existent files)
+            if not os.path.exists(abspath):
+                logger.info(f"Skipping non-existent file {abspath}")
+                continue
+            # Skip socket files
+            if stat.S_ISSOCK(os.stat(abspath).st_mode):
+                logger.info(f"Skip socket file {abspath}")
+                continue
+            if ignore_group:
+                if ignore_group.is_ignored(abspath):
+                    continue
+
+            ff.append(abspath)
+        all_files.extend(ff)
+
+        # Remove directories that we've already visited from dirnames
+        if deref_symlinks:
+            dirnames[:] = [d for d in dirnames if os.stat(os.path.join(root, d)).st_ino not in visited_inodes]
+
+    return all_files
+
+
+def _file_is_in_directory(file: str, directory: str) -> bool:
+    """Return True if file is in directory and in its children."""
+    try:
+        return os.path.commonpath([file, directory]) == directory
+    except ValueError as e:
+        # ValueError is raised by windows if the paths are not from the same drive
+        logger.debug(f"{file} and {directory} are not in the same drive: {e!s}")
+        return False
+
+
+def list_imported_modules_as_files(source_path: str, modules: List[ModuleType]) -> List[str]:
+    """Copies modules into destination that are in modules. The module files are copied only if:
+
+    1. Not a site-packages. These are installed packages and not user files.
+    2. Not in the sys.base_prefix or sys.prefix. These are also installed and not user files.
+    3. Does not share a common path with the source_path.
+    """
+    # source path is the folder holding the main script.
+    # but in register/package case, there are multiple folders.
+    # identify a common root amongst the packages listed?
+
+    import union
+    from union._utils.lazy_module import is_imported
+
+    files = []
+    union_root = os.path.dirname(union.__file__)
+
+    # These directories contain installed packages or modules from the Python standard library.
+    # If a module is from these directories, then they are not user files.
+    invalid_directories = [union_root, sys.prefix, sys.base_prefix, site.getusersitepackages(), *site.getsitepackages()]
+
+    for mod in modules:
+        # Be careful not to import a module with the .__file__ call if not yet imported.
+        if "LazyModule" in object.__getattribute__(mod, "__class__").__name__:
+            name = object.__getattribute__(mod, "__name__")
+            if is_imported(name):
+                mod_file = mod.__file__
+            else:
+                continue
+        else:
+            try:
+                mod_file = mod.__file__
+            except AttributeError:
+                continue
+
+        if mod_file is None:
+            continue
+
+        if any(_file_is_in_directory(mod_file, directory) for directory in invalid_directories):
+            continue
+
+        if not _file_is_in_directory(mod_file, source_path):
+            # Only upload files where the module file in the source directory
+            continue
+
+        files.append(mod_file)
+
+    return files
+
+
+def add_imported_modules_from_source(source_path: str, destination: str, modules: List[ModuleType]):
+    """Copies modules into destination that are in modules. The module files are copied only if:
+
+    1. Not a site-packages. These are installed packages and not user files.
+    2. Not in the sys.base_prefix or sys.prefix. These are also installed and not user files.
+    3. Does not share a common path with the source_path.
+    """
+    # source path is the folder holding the main script.
+    # but in register/package case, there are multiple folders.
+    # identify a common root amongst the packages listed?
+
+    files = list_imported_modules_as_files(source_path, modules)
+    for file in files:
+        relative_path = os.path.relpath(file, start=source_path)
+        new_destination = os.path.join(destination, relative_path)
+
+        if os.path.exists(new_destination):
+            # No need to copy if it already exists
+            continue
+
+        os.makedirs(os.path.dirname(new_destination), exist_ok=True)
+        shutil.copy(file, new_destination)
+
+
+def import_module_from_file(module_name, file):
+    try:
+        spec = importlib.util.spec_from_file_location(module_name, file)
+        module = importlib.util.module_from_spec(spec)
+        return module
+    except Exception as exc:
+        raise ModuleNotFoundError(f"Module from file {file} cannot be loaded") from exc
+
+
+def get_all_modules(source_path: str, module_name: Optional[str]) -> List[ModuleType]:
+    """Import python file with module_name in source_path and return all modules."""
+    sys_modules = list(sys.modules.values())
+    if module_name is None or module_name in sys.modules:
+        # module already exists, there is no need to import it again
+        return sys_modules
+
+    full_module = os.path.join(source_path, *module_name.split("."))
+    full_module_path = f"{full_module}.py"
+
+    is_python_file = os.path.exists(full_module_path) and os.path.isfile(full_module_path)
+    if not is_python_file:
+        return sys_modules
+
+    try:
+        new_module = import_module_from_file(module_name, full_module_path)
+        return [*sys_modules, new_module]
+    except Exception as exc:
+        logger.error(f"Using system modules, failed to import {module_name} from {full_module_path}: {exc!s}")
+        # Import failed so we fallback to `sys_modules`
+        return sys_modules
+
+
+@lru_cache
+def hash_file(file_path: typing.Union[os.PathLike, str]) -> (bytes, str, int):
+    """
+    Hash a file and produce a digest to be used as a version
+    """
+    h = hashlib.md5()
+    size = 0
+
+    with open(file_path, "rb") as file:
+        while True:
+            # Reading is buffered, so we can read smaller chunks.
+            chunk = file.read(h.block_size)
+            if not chunk:
+                break
+            h.update(chunk)
+            size += len(chunk)
+
+    return h.digest(), h.hexdigest(), size
+
+
+def _find_project_root(source_path) -> str:
+    """
+    Find the root of the project.
+    The root of the project is considered to be the first ancestor from source_path that does
+    not contain a __init__.py file.
+
+    N.B.: This assumption only holds for regular packages (as opposed to namespace packages)
+    """
+    # Start from the directory right above source_path
+    path = Path(source_path).parent.resolve()
+    while os.path.exists(os.path.join(path, "__init__.py")):
+        path = path.parent
+    return str(path)

union/_code_bundle/bundle.py

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

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