flyte 0.1.0__py3-none-any.whl → 0.2.0a0__py3-none-any.whl

flyte/_code_bundle/bundle.py

@@ -0,0 +1,209 @@
+import asyncio
+import gzip
+import logging
+import os
+import pathlib
+import tempfile
+from pathlib import Path
+from typing import ClassVar, Type
+
+from flyteidl.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, 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(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)
+                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: 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 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 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(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 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)
+        # Download the tgz file
+        path = await storage.get(bundle.tgz, str(downloaded_bundle.absolute()))
+        downloaded_bundle = pathlib.Path(path)
+        # 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
+        path = await storage.get(bundle.pkl, str(downloaded_bundle.absolute()))
+        downloaded_bundle = pathlib.Path(path)
+        return downloaded_bundle.absolute()
+    else:
+        raise ValueError("Code bundle should be either tgz or pkl, found neither.")

flyte/_context.py

@@ -0,0 +1,152 @@
+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:
+        """
+        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]:
+    """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, 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/_deploy.py

@@ -0,0 +1,243 @@
+from __future__ import annotations
+
+import asyncio
+from dataclasses import dataclass
+from typing import TYPE_CHECKING, Dict, List, Optional, Tuple
+
+import rich.repr
+
+import flyte.errors
+from flyte.models import SerializationContext
+from flyte.syncify import syncify
+
+from ._environment import Environment
+from ._image import Image
+from ._initialize import ensure_client, get_client, get_common_config, requires_initialization
+from ._logging import logger
+from ._task import TaskTemplate
+from ._task_environment import TaskEnvironment
+
+if TYPE_CHECKING:
+    from flyte._protos.workflow import task_definition_pb2
+
+    from ._code_bundle import CopyFiles
+    from ._internal.imagebuild.image_builder import ImageCache
+
+
+@rich.repr.auto
+@dataclass
+class DeploymentPlan:
+    envs: Dict[str, Environment]
+    version: Optional[str] = None
+
+
+@rich.repr.auto
+@dataclass
+class Deployment:
+    envs: Dict[str, Environment]
+    deployed_tasks: List[task_definition_pb2.TaskSpec] | None = None
+
+    def summary_repr(self) -> str:
+        """
+        Returns a summary representation of the deployment.
+        """
+        env_names = ", ".join(self.envs.keys())
+        task_names_versions = ", ".join(
+            f"{task.task_template.id.name} (v{task.task_template.id.version})" for task in self.deployed_tasks or []
+        )
+        return f"Deployment(envs=[{env_names}], tasks=[{task_names_versions}])"
+
+    def task_repr(self) -> List[List[Tuple[str, str]]]:
+        """
+        Returns a detailed representation of the deployed tasks.
+        """
+        tuples = []
+        if self.deployed_tasks:
+            for task in self.deployed_tasks:
+                tuples.append(
+                    [
+                        ("name", task.task_template.id.name),
+                        ("version", task.task_template.id.version),
+                    ]
+                )
+        return tuples
+
+    def env_repr(self) -> List[List[Tuple[str, str]]]:
+        """
+        Returns a detailed representation of the deployed environments.
+        """
+        tuples = []
+        for env_name, env in self.envs.items():
+            tuples.append(
+                [
+                    ("environment", env_name),
+                    ("image", env.image.uri if isinstance(env.image, Image) else env.image or ""),
+                ]
+            )
+        return tuples
+
+
+async def _deploy_task(
+    task: TaskTemplate, serialization_context: SerializationContext, dryrun: bool = False
+) -> task_definition_pb2.TaskSpec:
+    """
+    Deploy the given task.
+    """
+    ensure_client()
+    from ._internal.runtime.convert import convert_upload_default_inputs
+    from ._internal.runtime.task_serde import translate_task_to_wire
+    from ._protos.workflow import task_definition_pb2, task_service_pb2
+
+    image_uri = task.image.uri if isinstance(task.image, Image) else task.image
+
+    if dryrun:
+        return translate_task_to_wire(task, serialization_context)
+
+    default_inputs = await convert_upload_default_inputs(task.interface)
+    spec = translate_task_to_wire(task, serialization_context, default_inputs=default_inputs)
+
+    msg = f"Deploying task {task.name}, with image {image_uri} version {serialization_context.version}"
+    if spec.task_template.HasField("container") and spec.task_template.container.args:
+        msg += f" from {spec.task_template.container.args[-3]}.{spec.task_template.container.args[-1]}"
+    logger.info(msg)
+    task_id = task_definition_pb2.TaskIdentifier(
+        org=spec.task_template.id.org,
+        project=spec.task_template.id.project,
+        domain=spec.task_template.id.domain,
+        version=spec.task_template.id.version,
+        name=spec.task_template.id.name,
+    )
+
+    await get_client().task_service.DeployTask(task_service_pb2.DeployTaskRequest(task_id=task_id, spec=spec))
+    logger.info(f"Deployed task {task.name} with version {task_id.version}")
+    return spec
+
+
+async def _build_image_bg(env_name: str, image: Image) -> Tuple[str, str]:
+    """
+    Build the image in the background and return the environment name and the built image.
+    """
+    from ._build import build
+
+    logger.info(f"Building image {image.name} for environment {env_name}")
+    return env_name, await build.aio(image)
+
+
+async def build_images(deployment: DeploymentPlan) -> ImageCache:
+    """
+    Build the images for the given deployment plan and update the environment with the built image.
+    """
+    from ._internal.imagebuild.image_builder import ImageCache
+
+    images = []
+    image_identifier_map = {}
+    for env_name, env in deployment.envs.items():
+        if not isinstance(env.image, str):
+            logger.warning(f"Building Image for environment {env_name}, image: {env.image}")
+            images.append(_build_image_bg(env_name, env.image))
+
+        elif env.image == "auto" and "auto" not in image_identifier_map:
+            auto_image = Image.auto()
+            image_identifier_map["auto"] = auto_image.uri
+    final_images = await asyncio.gather(*images)
+
+    for env_name, image_uri in final_images:
+        logger.warning(f"Built Image for environment {env_name}, image: {image_uri}")
+        env = deployment.envs[env_name]
+        if isinstance(env.image, Image):
+            image_identifier_map[env.image.identifier] = env.image.uri
+
+    return ImageCache(image_lookup=image_identifier_map)
+
+
+@requires_initialization
+async def apply(deployment: DeploymentPlan, copy_style: CopyFiles, dryrun: bool = False) -> Deployment:
+    from ._code_bundle import build_code_bundle
+
+    cfg = get_common_config()
+    image_cache = await build_images(deployment)
+
+    version = deployment.version
+    code_bundle = None
+    if copy_style == "none" and not version:
+        raise flyte.errors.DeploymentError("Version must be set when copy_style is none")
+    else:
+        code_bundle = await build_code_bundle(from_dir=cfg.root_dir, dryrun=dryrun, copy_style=copy_style)
+        version = version or code_bundle.computed_version
+        # TODO we should update the version to include the image cache digest and code bundle digest. This is
+        # to ensure that changes in image dependencies, cause an update to the deployment version.
+        # TODO Also hash the environment and tasks to ensure that changes in the environment or tasks
+
+    sc = SerializationContext(
+        project=cfg.project,
+        domain=cfg.domain,
+        org=cfg.org,
+        code_bundle=code_bundle,
+        version=version,
+        image_cache=image_cache,
+        root_dir=cfg.root_dir,
+    )
+
+    tasks = []
+    for env_name, env in deployment.envs.items():
+        logger.info(f"Deploying environment {env_name}")
+        # TODO Make this pluggable based on the environment type
+        if isinstance(env, TaskEnvironment):
+            for task in env.tasks.values():
+                tasks.append(_deploy_task(task, dryrun=dryrun, serialization_context=sc))
+    return Deployment(envs=deployment.envs, deployed_tasks=await asyncio.gather(*tasks))
+
+
+def _recursive_discover(
+    planned_envs: Dict[str, Environment], envs: Environment | List[Environment]
+) -> Dict[str, Environment]:
+    """
+    Recursively deploy the environment and its dependencies, if not already deployed (present in env_tasks) and
+    return the updated env_tasks.
+    """
+    if isinstance(envs, Environment):
+        envs = [envs]
+    for env in envs:
+        # Skip if the environment is already planned
+        if env.name in planned_envs:
+            continue
+        # Recursively discover dependent environments
+        for dependent_env in env.depends_on:
+            _recursive_discover(planned_envs, dependent_env)
+        # Add the environment to the existing envs
+        planned_envs[env.name] = env
+    return planned_envs
+
+
+def plan_deploy(*envs: Environment, version: Optional[str] = None) -> DeploymentPlan:
+    if envs is None:
+        return DeploymentPlan({})
+    planned_envs = _recursive_discover({}, *envs)
+    return DeploymentPlan(planned_envs, version=version)
+
+
+@syncify
+async def deploy(
+    *envs: Environment,
+    dryrun: bool = False,
+    version: str | None = None,
+    interactive_mode: bool | None = None,
+    copy_style: CopyFiles = "loaded_modules",
+) -> Deployment:
+    """
+    Deploy the given environment or list of environments.
+    :param envs: Environment or list of environments to deploy.
+    :param dryrun: dryrun mode, if True, the deployment will not be applied to the control plane.
+    :param version: version of the deployment, if None, the version will be computed from the code bundle.
+    TODO: Support for interactive_mode
+    :param interactive_mode: Optional, can be forced to True or False.
+       If not provided, it will be set based on the current environment. For example Jupyter notebooks are considered
+         interactive mode, while scripts are not. This is used to determine how the code bundle is created.
+    :param copy_style: Copy style to use when running the task
+
+    :return: Deployment object containing the deployed environments and tasks.
+    """
+    if interactive_mode:
+        raise NotImplementedError("Interactive mode not yet implemented for deployment")
+    deployment = plan_deploy(*envs, version=version)
+    return await apply(deployment, copy_style=copy_style, dryrun=dryrun)

flyte/_doc.py

@@ -0,0 +1,29 @@
+import inspect
+from dataclasses import dataclass
+from typing import Callable, Optional
+
+
+@dataclass
+class Documentation:
+    """
+    This class is used to store the documentation of a task.
+
+    It can be set explicitly or extracted from the docstring of the task.
+    """
+
+    description: str
+
+    def __help__str__(self):
+        return self.description
+
+
+def extract_docstring(func: Optional[Callable]) -> Documentation:
+    """
+    Extracts the description from a docstring.
+    """
+    if func is None:
+        return Documentation(description="")
+    docstring = inspect.getdoc(func)
+    if not docstring:
+        return Documentation(description="")
+    return Documentation(description=docstring)

flyte/_docstring.py

@@ -0,0 +1,32 @@
+from typing import TYPE_CHECKING, Callable, Dict, Optional
+
+if TYPE_CHECKING:
+    pass
+
+
+class Docstring(object):
+    def __init__(self, docstring: Optional[str] = None, callable_: Optional[Callable] = None):
+        import docstring_parser
+
+        self._parsed_docstring: docstring_parser.Docstring
+
+        if docstring is not None:
+            self._parsed_docstring = docstring_parser.parse(docstring)
+        elif callable_.__doc__ is not None:
+            self._parsed_docstring = docstring_parser.parse(callable_.__doc__)
+
+    @property
+    def input_descriptions(self) -> Dict[str, Optional[str]]:
+        return {p.arg_name: p.description for p in self._parsed_docstring.params}
+
+    @property
+    def output_descriptions(self) -> Dict[str, Optional[str]]:
+        return {p.return_name: p.description for p in self._parsed_docstring.many_returns if p.return_name is not None}
+
+    @property
+    def short_description(self) -> Optional[str]:
+        return self._parsed_docstring.short_description
+
+    @property
+    def long_description(self) -> Optional[str]:
+        return self._parsed_docstring.long_description