flyte 0.0.1b0__py3-none-any.whl

flyte/_code_bundle/_packaging.py

@@ -0,0 +1,187 @@
+from __future__ import annotations
+
+import gzip
+import hashlib
+import os
+import pathlib
+import posixpath
+import shutil
+import stat
+import subprocess
+import tarfile
+import time
+import typing
+from typing import List, Optional, Tuple, Union
+
+import click
+from rich import print as rich_print
+from rich.tree import Tree
+
+from flyte._logging import logger
+
+from ._ignore import Ignore, IgnoreGroup
+from ._utils import CopyFiles, _filehash_update, _pathhash_update, ls_files, tar_strip_file_attributes
+
+FAST_PREFIX = "fast"
+FAST_FILEENDING = ".tar.gz"
+
+
+def print_ls_tree(source: os.PathLike, ls: typing.List[str]):
+    click.secho("Files to be copied for fast registration...", fg="bright_blue")
+
+    tree_root = Tree(
+        f":open_file_folder: [link file://{source}]{source} (detected source root)",
+        guide_style="bold bright_blue",
+    )
+    trees = {pathlib.Path(source): tree_root}
+
+    for f in ls:
+        fpp = pathlib.Path(f)
+        if fpp.parent not in trees:
+            # add trees for all intermediate folders
+            current = tree_root
+            current_path = pathlib.Path(source)
+            for subdir in fpp.parent.relative_to(source).parts:
+                current_path = current_path / subdir
+                if current_path not in trees:
+                    current = current.add(f"{subdir}", guide_style="bold bright_blue")
+                    trees[current_path] = current
+                else:
+                    current = trees[current_path]
+        trees[fpp.parent].add(f"{fpp.name}", guide_style="bold bright_blue")
+    rich_print(tree_root)
+
+
+def _compress_tarball(source: pathlib.Path, output: pathlib.Path) -> None:
+    """Compress code tarball using pigz if available, otherwise gzip"""
+    if pigz := shutil.which("pigz"):
+        with open(str(output), "wb") as gzipped:
+            subprocess.run([pigz, "--no-time", "-c", str(source)], stdout=gzipped, check=True)
+    else:
+        start_time = time.time()
+        with gzip.GzipFile(filename=str(output), mode="wb", mtime=0) as gzipped:
+            with open(source, "rb") as source_file:
+                gzipped.write(source_file.read())
+
+        end_time = time.time()
+        warning_time = 10
+        if end_time - start_time > warning_time:
+            click.secho(
+                f"Code tarball compression took {end_time - start_time:.0f} seconds. "
+                f"Consider installing `pigz` for faster compression.",
+                fg="yellow",
+            )
+
+
+def list_files_to_bundle(
+    source: pathlib.Path,
+    deref_symlinks: bool = False,
+    *ignores: typing.Type[Ignore],
+    copy_style: CopyFiles = "all",
+) -> typing.Tuple[List[str], str]:
+    """
+    Takes a source directory and returns a list of all files to be included in the code bundle and a hexdigest of the
+    included files.
+    :param source: The source directory to package
+    :param deref_symlinks: Whether to dereference symlinks or not
+    :param ignores: A list of Ignore classes to use for ignoring files
+    :param copy_style: The copy style to use for the tarball
+    :return: A list of all files to be included in the code bundle and a hexdigest of the included files
+    """
+    ignore = IgnoreGroup(source, *ignores)
+
+    ls, ls_digest = ls_files(source, copy_style, deref_symlinks, ignore)
+    logger.debug(f"Hash digest: {ls_digest}")
+    return ls, ls_digest
+
+
+def create_bundle(
+    source: pathlib.Path, output_dir: pathlib.Path, ls: List[str], ls_digest: str, deref_symlinks: bool = False
+) -> Tuple[pathlib.Path, float, float]:
+    """
+    Takes a source directory and packages everything not covered by common ignores into a tarball.
+    The output_dir is the directory where the tarball and a compressed version of the tarball will be written.
+    The output_dir can be a temporary directory.
+
+    :param source: The source directory to package
+    :param output_dir: The directory to write the tarball to
+    :param deref_symlinks: Whether to dereference symlinks or not
+    :param ls: The list of files to include in the tarball
+    :param ls_digest: The hexdigest of the included files
+    :return: The path to the tarball, the size of the tarball in MB, and the size of the compressed tarball in MB
+    """
+    # Compute where the archive should be written
+    archive_fname = output_dir / f"{FAST_PREFIX}{ls_digest}{FAST_FILEENDING}"
+    tar_path = output_dir / "tmp.tar"
+    with tarfile.open(str(tar_path), "w", dereference=deref_symlinks) as tar:
+        for ws_file in ls:
+            rel_path = os.path.relpath(ws_file, start=source)
+            tar.add(
+                os.path.join(source, ws_file),
+                recursive=False,
+                arcname=rel_path,
+                filter=lambda x: tar_strip_file_attributes(x),
+            )
+
+    size_mbs = tar_path.stat().st_size / 1024 / 1024
+    _compress_tarball(tar_path, archive_fname)
+    asize_mbs = archive_fname.stat().st_size / 1024 / 1024
+
+    return archive_fname, size_mbs, asize_mbs
+
+
+def compute_digest(source: Union[os.PathLike, List[os.PathLike]], filter: Optional[typing.Callable] = None) -> str:
+    """
+    Walks the entirety of the source dir to compute a deterministic md5 hex digest of the dir contents.
+    :param os.PathLike source:
+    :param callable filter:
+    :return Text:
+    """
+    hasher = hashlib.md5()
+
+    def compute_digest_for_file(path: os.PathLike, rel_path: os.PathLike) -> None:
+        # Only consider files that exist (e.g. disregard symlinks that point to non-existent files)
+        if not os.path.exists(path):
+            logger.info(f"Skipping non-existent file {path}")
+            return
+
+        # Skip socket files
+        if stat.S_ISSOCK(os.stat(path).st_mode):
+            logger.info(f"Skip socket file {path}")
+            return
+
+        if filter:
+            if filter(rel_path):
+                return
+
+        _filehash_update(path, hasher)
+        _pathhash_update(rel_path, hasher)
+
+    def compute_digest_for_dir(source: os.PathLike) -> None:
+        for root, _, files in os.walk(str(source), topdown=True):
+            files.sort()
+
+            for fname in files:
+                abspath = os.path.join(root, fname)
+                relpath = os.path.relpath(abspath, source)
+                compute_digest_for_file(pathlib.Path(abspath), pathlib.Path(relpath))
+
+    if isinstance(source, list):
+        for src in source:
+            if os.path.isdir(src):
+                compute_digest_for_dir(src)
+            else:
+                compute_digest_for_file(src, os.path.basename(src))
+    else:
+        compute_digest_for_dir(source)
+
+    return hasher.hexdigest()
+
+
+def get_additional_distribution_loc(remote_location: str, identifier: str) -> str:
+    """
+    :param Text remote_location:
+    :param Text identifier:
+    :return Text:
+    """
+    return posixpath.join(remote_location, "{}.{}".format(identifier, "tar.gz"))

flyte/_code_bundle/_utils.py

@@ -0,0 +1,339 @@
+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 flyte._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 command with. For flyte 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 = str(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(str(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. Shares a common path with the source_path.
+    """
+
+    import flyte
+    from flyte._utils.lazy_module import is_imported
+
+    files = []
+    union_root = os.path.dirname(flyte.__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]) -> Tuple[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)