flyte 2.0.0b32__py3-none-any.whl

flyte/_code_bundle/bundle.py

@@ -0,0 +1,266 @@
+import asyncio
+import gzip
+import logging
+import os
+import pathlib
+import tempfile
+from pathlib import Path
+from typing import ClassVar, Type
+
+from async_lru import alru_cache
+from flyteidl2.core.tasks_pb2 import TaskTemplate
+
+from flyte._logging import log, logger
+from flyte._utils import AsyncLRUCache
+from flyte.models import CodeBundle
+
+from ._ignore import GitIgnore, Ignore, StandardIgnore
+from ._packaging import create_bundle, list_files_to_bundle, list_relative_files_to_bundle, print_ls_tree
+from ._utils import CopyFiles, hash_file
+
+_pickled_file_extension = ".pkl.gz"
+_tar_file_extension = ".tar.gz"
+
+
+class _PklCache:
+    _pkl_cache: ClassVar[AsyncLRUCache[str, str]] = AsyncLRUCache[str, str](maxsize=100)
+
+    @classmethod
+    async def put(cls, digest: str, upload_to_path: str, from_path: pathlib.Path) -> str:
+        """
+        Get the pickled code bundle from the cache or build it if not present.
+
+        :param digest: The hash digest of the task template.
+        :param upload_to_path: The path to upload the pickled file to.
+        :param from_path: The path to read the pickled file from.
+        :return: CodeBundle object containing the pickled file path and the computed version.
+        """
+        import flyte.storage as storage
+
+        async def put_data() -> str:
+            return await storage.put(str(from_path), to_path=str(upload_to_path))
+
+        return await cls._pkl_cache.get(
+            key=digest,
+            value_func=put_data,
+        )
+
+
+async def build_pkl_bundle(
+    o: TaskTemplate,
+    upload_to_controlplane: bool = True,
+    upload_from_dataplane_base_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_base_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
+
+    if upload_to_controlplane and upload_from_dataplane_base_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 flyte.remote import upload_file
+
+            hash_digest, remote_path = await upload_file.aio(dest)
+            return CodeBundle(pkl=remote_path, computed_version=hash_digest)
+
+        elif upload_from_dataplane_base_path:
+            from flyte._internal.runtime import io
+
+            _, str_digest, _ = hash_file(file_path=dest)
+            upload_path = io.pkl_path(upload_from_dataplane_base_path, str_digest)
+            logger.debug(f"Uploading pickled code bundle to dataplane path {upload_path}.")
+            final_path = await _PklCache.put(
+                digest=str_digest,
+                upload_to_path=upload_path,
+                from_path=dest,
+            )
+            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, follow_symlinks=True)
+                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)
+
+
+@alru_cache
+async def build_code_bundle(
+    from_dir: Path,
+    *ignore: Type[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 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 flyte.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.aio(bundle_path)
+            logger.debug(f"Code bundle uploaded to {remote_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)
+
+
+@alru_cache
+async def build_code_bundle_from_relative_paths(
+    relative_paths: tuple[str, ...],
+    from_dir: Path,
+    extract_dir: str = ".",
+    dryrun: bool = False,
+    copy_bundle_to: pathlib.Path | None = None,
+) -> CodeBundle:
+    """
+    Build a code bundle from a list of relative paths.
+    :param relative_paths: The list of relative paths to bundle.
+    :param from_dir: The directory of the code to bundle. This is the root directory for the source.
+    :param ignore: The list of ignores to apply. This is a list of Ignore classes.
+    :param extract_dir: The directory to extract the code bundle to, when in the container. It defaults to the current
+        working directory.
+    :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 relative paths.")
+    from flyte.remote import upload_file
+
+    logger.debug("Finding files to bundle")
+    files, digest = list_relative_files_to_bundle(relative_paths, from_dir)
+    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.aio(bundle_path)
+            logger.debug(f"Code bundle uploaded to {remote_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.
+    """
+    import sys
+
+    import flyte.storage as storage
+
+    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)
+        if downloaded_bundle.exists():
+            return downloaded_bundle.absolute()
+        # Download the tgz file
+        await storage.get(bundle.tgz, str(downloaded_bundle.absolute()))
+        # 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.
+        args = [
+            "-xvf",
+            str(downloaded_bundle),
+            "-C",
+            str(dest),
+        ]
+        if sys.platform != "darwin":
+            args.insert(0, "--overwrite")
+
+        process = await asyncio.create_subprocess_exec(
+            "tar",
+            *args,
+            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
+        await storage.get(bundle.pkl, str(downloaded_bundle.absolute()))
+        return downloaded_bundle.absolute()
+    else:
+        raise ValueError("Code bundle should be either tgz or pkl, found neither.")

flyte/_constants.py

@@ -0,0 +1 @@
+FLYTE_SYS_PATH = "_F_SYS_PATH"  # The paths that will be appended to sys.path at runtime

flyte/_context.py

@@ -0,0 +1,155 @@
+from __future__ import annotations
+
+import contextvars
+from dataclasses import dataclass, replace
+from typing import TYPE_CHECKING, Awaitable, Callable, Optional, ParamSpec, TypeVar
+
+from flyte._logging import logger
+from flyte.models import GroupData, RawDataPath, TaskContext
+
+if TYPE_CHECKING:
+    from flyte.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
+        self._token = None  # Context variable token to restore the previous context
+
+    @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: bool
+        """
+        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."""
+        try:
+            root_context_var.reset(self._token)
+        except Exception as e:
+            logger.warn(f"Failed to reset context: {e}")
+            raise e
+
+    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]:
+    """
+    Returns flyte.models.TaskContext if within a task context, else None
+    Note: Only use this in task code and not module level.
+    """
+    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, Awaitable[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)

flyte/_custom_context.py

@@ -0,0 +1,73 @@
+from __future__ import annotations
+
+from contextlib import contextmanager
+
+from flyte._context import ctx
+
+from ._context import internal_ctx
+
+
+def get_custom_context() -> dict[str, str]:
+    """
+    Get the current input context. This can be used within a task to retrieve
+    context metadata that was passed to the action.
+
+    Context will automatically propagate to sub-actions.
+
+    Example:
+    ```python
+    import flyte
+
+    env = flyte.TaskEnvironment(name="...")
+
+    @env.task
+    def t1():
+        # context can be retrieved with `get_custom_context`
+        ctx = flyte.get_custom_context()
+        print(ctx)  # {'project': '...', 'entity': '...'}
+    ```
+
+    :return: Dictionary of context key-value pairs
+    """
+    tctx = ctx()
+    if tctx is None or tctx.custom_context is None:
+        return {}
+    return tctx.custom_context
+
+
+@contextmanager
+def custom_context(**context: str):
+    """
+    Synchronous context manager to set input context for tasks spawned within this block.
+
+    Example:
+    ```python
+    import flyte
+
+    env = flyte.TaskEnvironment(name="...")
+
+    @env.task
+    def t1():
+        ctx = flyte.get_custom_context()
+        print(ctx)
+
+    @env.task
+    def main():
+        # context can be passed via a context manager
+        with flyte.custom_context(project="my-project"):
+            t1()  # will have {'project': 'my-project'} as context
+    ```
+
+    :param context: Key-value pairs to set as input context
+    """
+    ctx = internal_ctx()
+    if ctx.data.task_context is None:
+        yield
+        return
+
+    tctx = ctx.data.task_context
+    new_tctx = tctx.replace(custom_context={**tctx.custom_context, **context})
+
+    with ctx.replace_task_context(new_tctx):
+        yield
+        # Exit the context and restore the previous context

flyte/_debug/constants.py

@@ -0,0 +1,38 @@
+import os
+from pathlib import Path
+
+# Where the code-server tar and plugins are downloaded to
+EXECUTABLE_NAME = "code-server"
+DOWNLOAD_DIR = Path.home() / ".code-server"
+HOURS_TO_SECONDS = 60 * 60
+DEFAULT_UP_SECONDS = 10 * HOURS_TO_SECONDS  # 10 hours
+DEFAULT_CODE_SERVER_REMOTE_PATHS = {
+    "amd64": "https://github.com/coder/code-server/releases/download/v4.18.0/code-server-4.18.0-linux-amd64.tar.gz",
+    "arm64": "https://github.com/coder/code-server/releases/download/v4.18.0/code-server-4.18.0-linux-arm64.tar.gz",
+}
+DEFAULT_CODE_SERVER_EXTENSIONS = [
+    "https://raw.githubusercontent.com/flyteorg/flytetools/master/flytekitplugins/flyin/ms-python.python-2023.20.0.vsix",
+]
+
+# Duration to pause the checking of the heartbeat file until the next one
+HEARTBEAT_CHECK_SECONDS = 60
+MAX_IDLE_SECONDS = 180
+
+# The path is hardcoded by code-server
+# https://coder.com/docs/code-server/latest/FAQ#what-is-the-heartbeat-file
+HEARTBEAT_PATH = os.path.expanduser("~/.local/share/code-server/heartbeat")
+
+INTERACTIVE_DEBUGGING_FILE_NAME = "flyteinteractive_interactive_entrypoint.py"
+RESUME_TASK_FILE_NAME = "flyteinteractive_resume_task.py"
+# Config keys to store in task template
+VSCODE_TYPE_KEY = "flyteinteractive_type"
+VSCODE_PORT_KEY = "flyteinteractive_port"
+
+TASK_FUNCTION_SOURCE_PATH = "TASK_FUNCTION_SOURCE_PATH"
+
+# Default max idle seconds to terminate the flyteinteractive server
+HOURS_TO_SECONDS = 60 * 60
+MAX_IDLE_SECONDS = 10 * HOURS_TO_SECONDS  # 10 hours
+
+# Subprocess constants
+EXIT_CODE_SUCCESS = 0

flyte/_debug/utils.py

@@ -0,0 +1,17 @@
+import asyncio
+
+from flyte._debug.constants import EXIT_CODE_SUCCESS
+from flyte._logging import logger
+
+
+async def execute_command(cmd: str):
+    """
+    Execute a command in the shell.
+    """
+    process = await asyncio.create_subprocess_shell(cmd, stdout=asyncio.subprocess.PIPE, stderr=asyncio.subprocess.PIPE)
+    logger.info(f"cmd: {cmd}")
+    stdout, stderr = await process.communicate()
+    if process.returncode != EXIT_CODE_SUCCESS:
+        raise RuntimeError(f"Command {cmd} failed with error: {stderr!r}")
+    logger.info(f"stdout: {stdout!r}")
+    logger.info(f"stderr: {stderr!r}")