flyte 2.0.0b32__py3-none-any.whl

flyte/_code_bundle/_packaging.py

@@ -0,0 +1,218 @@
+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.tree import Tree
+
+from flyte._logging import _get_console, logger
+
+from ._ignore import Ignore, IgnoreGroup
+from ._utils import (
+    CopyFiles,
+    _filehash_update,
+    _pathhash_update,
+    ls_files,
+    ls_relative_files,
+    tar_strip_file_attributes,
+)
+
+FAST_PREFIX = "fast"
+FAST_FILEENDING = ".tar.gz"
+
+
+def print_ls_tree(source: os.PathLike, ls: typing.List[str]):
+    logger.info("Files to be copied for fast registration...")
+
+    tree_root = Tree(
+        f"File structure:\n:open_file_folder: {source}",
+        guide_style="bold bright_blue",
+    )
+    source_path = pathlib.Path(source).resolve()
+    trees = {source_path: 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 = source_path  # pathlib.Path(source)
+            for subdir in fpp.parent.relative_to(source_path).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")
+
+    console = _get_console()
+    with console.capture() as capture:
+        console.print(tree_root, overflow="ignore", no_wrap=True, crop=False)
+    logger.info(f"Root directory: [link=file://{source}]{source}[/link]")
+    logger.info(capture.get(), extra={"console": console})
+
+
+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 of files to be included in the code bundle: {ls_digest}")
+    return ls, ls_digest
+
+
+def list_relative_files_to_bundle(
+    relative_paths: tuple[str, ...],
+    source: pathlib.Path,
+) -> typing.Tuple[List[str], str]:
+    """
+    List the files in the relative paths.
+
+    :param relative_paths: The list of relative paths to bundle.
+    :param source: The source directory to package.
+    :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.
+    """
+    _source = source
+
+    all_files, digest = ls_relative_files(list(relative_paths), source)
+    logger.debug(f"Hash of files to be included in the code bundle: {digest}")
+    return all_files, 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,347 @@
+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 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 ls_relative_files(relative_paths: list[str], source_path: pathlib.Path) -> tuple[list[str], str]:
+    relative_paths = list(relative_paths)
+    relative_paths.sort()
+    hasher = hashlib.md5()
+
+    all_files = []
+    for file in relative_paths:
+        path = source_path / file
+        all_files.append(str(path))
+        _filehash_update(path, hasher)
+        _pathhash_update(path, 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 os.walk(source_path, topdown=True, followlinks=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 = (pathlib.Path(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 pathlib.Path(file).resolve().is_relative_to(pathlib.Path(directory).resolve())
+    except OSError as e:
+        # OSError can be raised if paths cannot be resolved (permissions, broken symlinks, etc.)
+        logger.debug(f"Failed to resolve paths for {file} and {directory}: {e!s}")
+        return False
+
+
+def list_imported_modules_as_files(source_path: str, modules: List[ModuleType]) -> List[str]:
+    """Lists the files of modules that have been loaded.  The files are only included 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 = set()
+    flyte_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 = [flyte_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
+            # print log line for files that have common ancestor with source_path, but not in it.
+            logger.debug(f"{mod_file} is not in {source_path}")
+            continue
+
+        if not pathlib.Path(mod_file).is_file():
+            # Some modules have a __file__ attribute that are relative to the base package. Let's skip these,
+            # can add more rigorous logic to really pull out the correct file location if we need to.
+            logger.debug(f"Skipping {mod_file} from {mod.__name__} because it is not a file")
+            continue
+
+        files.add(mod_file)
+
+    return list(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