flyte 0.0.1b0__py3-none-any.whl

union/_cli/_run.py

@@ -0,0 +1,150 @@
+from __future__ import annotations
+
+from union._logging import logger
+from dataclasses import dataclass, fields
+from pathlib import Path
+from types import ModuleType
+from typing import Any, Dict, List, cast
+
+import click
+
+import union
+
+from .._code_bundle._utils import CopyFiles
+from .._initialize import initialize_in_cluster, S3
+from .._task import TaskTemplate
+from . import _common as common
+from ._common import CLIConfig
+
+
+@dataclass
+class RunArguments:
+    project: str = common.make_click_option_field(common.PROJECT_OPTION)
+    domain: str = common.make_click_option_field(common.DOMAIN_OPTION)
+    local: bool = common.make_click_option_field(
+        click.Option(
+            ["--local"],
+            is_flag=True,
+            help="Run the task locally",
+        )
+    )
+    hybrid: bool = common.make_click_option_field(
+        click.Option(
+            ["--hybrid"],
+            is_flag=True,
+            help="hybrid mode",
+        )
+    )
+    copy_style: CopyFiles = common.make_click_option_field(
+        click.Option(
+            ["--copy-style"],
+            type=click.Choice(["loaded_modules", "all_files", "none"]),
+            default="loaded_modules",
+            help="Copy style to use when running the task",
+        )
+    )
+
+    @classmethod
+    def from_dict(cls, d: Dict[str, Any]) -> RunArguments:
+        return cls(**d)
+
+    @classmethod
+    def options(cls) -> List[click.Option]:
+        """
+        Return the set of base parameters added to run subcommand.
+        """
+        return [common.get_option_from_metadata(f.metadata) for f in fields(cls) if f.metadata]
+
+
+class RunTaskCommand(click.Command):
+    def __init__(self, obj_name: str, obj: Any, run_args: RunArguments, *args, **kwargs):
+        self.obj_name = obj_name
+        self.obj = cast(TaskTemplate, obj)
+        self.run_args = run_args
+        super().__init__(name=obj_name, *args, **kwargs)
+
+    def invoke(self, ctx):
+        obj: CLIConfig = ctx.obj
+        obj.init(self.run_args.project, self.run_args.domain)
+
+        # todo: remove this when backend supports
+        from union.storage import get_random_local_path
+
+        run_name = "random-run-" + str(get_random_local_path())[-8:]
+        if self.run_args.local:
+            mode = "local"
+        elif self.run_args.hybrid:
+            mode = "hybrid"
+            # TODO configure storage correctly for cluster
+            initialize_in_cluster(storage=S3.auto())
+        else:
+            mode = "remote"
+        logger.debug(f"Running {self.obj_name} in {mode} mode")
+        r = union.with_runcontext(
+            name=run_name,
+            copy_style=self.run_args.copy_style,
+            version=self.run_args.copy_style,
+            mode=mode,
+        ).run(self.obj)
+        # click.secho(f"Created Run: {r.action.id}", fg="green")
+
+
+class TaskPerFileGroup(common.ObjectsPerFileGroup):
+    """
+    Group that creates a command for each task in the current directory that is not __init__.py.
+    """
+
+    def __init__(self, filename: Path, run_args: RunArguments, *args, **kwargs):
+        super().__init__(*args, filename=filename, **kwargs)
+        self.run_args = run_args
+
+    def _filter_objects(self, module: ModuleType) -> Dict[str, Any]:
+        return {k: v for k, v in module.__dict__.items() if isinstance(v, TaskTemplate)}
+
+    def _get_command_for_obj(self, ctx: click.Context, obj_name: str, obj: Any) -> click.Command:
+        obj = cast(TaskTemplate, obj)
+        return RunTaskCommand(
+            obj_name=obj_name,
+            obj=obj,
+            help=obj.docs.__help__str__() if obj.docs else None,
+            run_args=self.run_args,
+        )
+
+
+class TaskFiles(common.FileGroup):
+    """
+    Group that creates a command for each file in the current directory that is not __init__.py.
+    """
+
+    common_options_enabled = False
+
+    def __init__(
+        self,
+        *args,
+        directory: Path | None = None,
+        **kwargs,
+    ):
+        if "params" not in kwargs:
+            kwargs["params"] = []
+        kwargs["params"].extend(RunArguments.options())
+        super().__init__(*args, directory=directory, **kwargs)
+
+    def get_command(self, ctx, filename):
+        run_args = RunArguments.from_dict(ctx.params)
+        fp = Path(filename)
+        if not fp.exists():
+            raise click.BadParameter(f"File {filename} does not exist")
+        if fp.is_dir():
+            return TaskFiles(directory=fp)
+        return TaskPerFileGroup(
+            filename=Path(filename),
+            run_args=run_args,
+            name=filename,
+            help=f"Run, functions decorated `env.task` or instances of Tasks in {filename}",
+        )
+
+
+run = TaskFiles(
+    name="run",
+    help="Run a task from a python file",
+)

union/_cli/main.py

@@ -0,0 +1,72 @@
+import rich_click as click
+
+from ._common import CLIConfig
+from ._create import create
+from ._deploy import deploy
+from ._get import get
+from ._run import run
+
+
+def _verbosity_to_loglevel(verbosity: int) -> int | None:
+    """
+    Converts a verbosity level from the CLI to a logging level.
+
+    :param verbosity: verbosity level from the CLI
+    :return: logging level
+    """
+    import logging
+
+    match verbosity:
+        case 0:
+            return None
+        case 1:
+            return logging.WARNING
+        case 2:
+            return logging.INFO
+        case _:
+            return logging.DEBUG
+
+
+@click.group(cls=click.RichGroup)
+@click.option(
+    "--endpoint",
+    type=str,
+    required=False,
+    help="The endpoint to connect to, this will override any config and simply used pkce to connect.",
+)
+@click.option(
+    "--insecure",
+    is_flag=True,
+    required=False,
+    help="insecure",
+    type=bool,
+    default=False,
+)
+@click.option(
+    "-v",
+    "--verbose",
+    required=False,
+    help="Show verbose messages and exception traces",
+    count=True,
+    default=0,
+    type=int,
+)
+@click.option(
+    "--org-override",
+    type=str,
+    required=False,
+    help="Override for org",
+)
+@click.pass_context
+def main(ctx: click.Context, endpoint: str | None, insecure: bool, verbose: int, org_override: str | None):
+    """
+    v2 cli. Root command, please use one of the subcommands.
+    """
+    log_level = _verbosity_to_loglevel(verbose)
+    ctx.obj = CLIConfig(log_level=log_level, endpoint=endpoint, insecure=insecure, org_override=org_override)
+
+
+main.add_command(run)
+main.add_command(deploy)
+main.add_command(get)  # type: ignore
+main.add_command(create)  # type: ignore

union/_code_bundle/__init__.py

@@ -0,0 +1,8 @@
+from ._ignore import GitIgnore, IgnoreGroup, StandardIgnore
+from ._utils import CopyFiles
+from .bundle import build_code_bundle, build_pkl_bundle, download_bundle
+
+__all__ = ["CopyFiles", "build_code_bundle", "build_pkl_bundle", "default_ignores", "download_bundle"]
+
+
+default_ignores = [GitIgnore, StandardIgnore, IgnoreGroup]

union/_code_bundle/_ignore.py

@@ -0,0 +1,113 @@
+import os
+import pathlib
+import subprocess
+import tarfile as _tarfile
+from abc import ABC, abstractmethod
+from fnmatch import fnmatch
+from pathlib import Path
+from shutil import which
+from typing import List, Optional, Type
+
+from union._logging import logger
+
+
+class Ignore(ABC):
+    """Base for Ignores, implements core logic. Children have to implement _is_ignored"""
+
+    def __init__(self, root: Path):
+        self.root = root
+
+    def is_ignored(self, path: pathlib.Path) -> bool:
+        return self._is_ignored(path)
+
+    def tar_filter(self, tarinfo: _tarfile.TarInfo) -> Optional[_tarfile.TarInfo]:
+        if self.is_ignored(pathlib.Path(tarinfo.name)):
+            return None
+        return tarinfo
+
+    @abstractmethod
+    def _is_ignored(self, path: pathlib.Path) -> bool:
+        pass
+
+
+class GitIgnore(Ignore):
+    """Uses git cli (if available) to list all ignored files and compare with those."""
+
+    def __init__(self, root: Path):
+        super().__init__(root)
+        self.has_git = which("git") is not None
+        self.ignored_files = self._list_ignored_files()
+        self.ignored_dirs = self._list_ignored_dirs()
+
+    def _git_wrapper(self, extra_args: List[str]) -> set[str]:
+        if self.has_git:
+            out = subprocess.run(
+                ["git", "ls-files", "-io", "--exclude-standard", *extra_args],
+                cwd=self.root,
+                capture_output=True,
+                check=False,
+            )
+            if out.returncode == 0:
+                return set(out.stdout.decode("utf-8").split("\n")[:-1])
+            logger.info(f"Could not determine ignored paths due to:\n{out.stderr}\nNot applying any filters")
+            return set()
+        logger.info("No git executable found, not applying any filters")
+        return set()
+
+    def _list_ignored_files(self) -> set[str]:
+        return self._git_wrapper([])
+
+    def _list_ignored_dirs(self) -> set[str]:
+        return self._git_wrapper(["--directory"])
+
+    def _is_ignored(self, path: str) -> bool:
+        if self.ignored_files:
+            # git-ls-files uses POSIX paths
+            if Path(path).as_posix() in self.ignored_files:
+                return True
+            # Ignore empty directories
+            if os.path.isdir(os.path.join(self.root, path)) and self.ignored_dirs:
+                return Path(path).as_posix() + "/" in self.ignored_dirs
+        return False
+
+
+STANDARD_IGNORE_PATTERNS = ["*.pyc", ".cache", ".cache/*", "__pycache__", "**/__pycache__"]
+
+
+class StandardIgnore(Ignore):
+    """Retains the standard ignore functionality that previously existed. Could in theory
+    by fed with custom ignore patterns from cli."""
+
+    def __init__(self, root: Path, patterns: Optional[List[str]] = None):
+        super().__init__(root)
+        self.patterns = patterns if patterns else STANDARD_IGNORE_PATTERNS
+
+    def _is_ignored(self, path: pathlib.Path) -> bool:
+        for pattern in self.patterns:
+            if fnmatch(str(path), pattern):
+                return True
+        return False
+
+
+class IgnoreGroup(Ignore):
+    """Groups multiple Ignores and checks a path against them. A file is ignored if any
+    Ignore considers it ignored."""
+
+    def __init__(self, root: Path, *ignores: Type[Ignore]):
+        super().__init__(root)
+        self.ignores = [ignore(root) for ignore in ignores]
+
+    def _is_ignored(self, path: pathlib.Path) -> bool:
+        for ignore in self.ignores:
+            if ignore.is_ignored(path):
+                return True
+        return False
+
+    def list_ignored(self) -> List[str]:
+        ignored = []
+        for dir, _, files in self.root.walk():
+            for file in files:
+                abs_path = dir / file
+                if self.is_ignored(abs_path):
+                    ignored.append(str(abs_path.relative_to(self.root)))
+        return ignored

union/_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 union._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[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(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(abspath, 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"))