flyte 0.0.1b0__py3-none-any.whl

flyte/_code_bundle/bundle.py

@@ -0,0 +1,178 @@
+import asyncio
+import gzip
+import logging
+import os
+import pathlib
+import tempfile
+from pathlib import Path
+from typing import Type
+
+from flyteidl.core.tasks_pb2 import TaskTemplate
+
+import flyte.storage as storage
+from flyte._datastructures import CodeBundle
+from flyte._logging import log, logger
+
+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"
+
+
+async def build_pkl_bundle(
+    o: TaskTemplate,
+    upload_to_controlplane: bool = True,
+    upload_from_dataplane_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_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
+
+    import flyte.storage as storage
+
+    if upload_to_controlplane and upload_from_dataplane_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_path:
+            logger.debug(f"Uploading pickled code bundle to dataplane path {upload_from_dataplane_path}.")
+            _, str_digest, _ = hash_file(file_path=dest)
+            final_path = await storage.put(str(dest), upload_from_dataplane_path)
+            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.info(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.
+    """
+    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,146 @@
+from __future__ import annotations
+
+import contextvars
+from dataclasses import dataclass, replace
+from typing import TYPE_CHECKING, Awaitable, Callable, Optional, ParamSpec, TypeVar
+
+from flyte._datastructures 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
+
+    @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."""
+        root_context_var.reset(self._token)
+
+    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/_datastructures.py

@@ -0,0 +1,342 @@
+from __future__ import annotations
+
+import inspect
+import os
+import pathlib
+import tempfile
+from dataclasses import dataclass, field, replace
+from typing import TYPE_CHECKING, Any, Callable, Dict, Optional, Tuple, Type
+
+from flyte._docstring import Docstring
+from flyte._interface import extract_return_annotation
+from flyte._logging import logger
+from flyte._utils.helpers import base36_encode
+
+if TYPE_CHECKING:
+    from flyte._internal.imagebuild.image_builder import ImageCache
+    from flyte.report import Report
+
+
+def generate_random_name() -> str:
+    """
+    Generate a random name for the task. This is used to create unique names for tasks.
+    TODO we can use unique-namer in the future, for now its just guids
+    """
+    from uuid import uuid4
+
+    return str(uuid4())  # Placeholder for actual random name generation logic
+
+
+@dataclass(frozen=True, kw_only=True)
+class ActionID:
+    """
+    A class representing the ID of an Action, nested within a Run. This is used to identify a specific action on a task.
+    """
+
+    name: str
+    run_name: str | None = None
+    project: str | None = None
+    domain: str | None = None
+    org: str | None = None
+
+    def __post_init__(self):
+        if self.run_name is None:
+            object.__setattr__(self, "run_name", self.name)
+
+    @classmethod
+    def create_random(cls):
+        name = generate_random_name()
+        return cls(name=name, run_name=name)
+
+    def new_sub_action(self, name: str | None = None) -> ActionID:
+        """
+        Create a new sub-run with the given name. If  name is None, a random name will be generated.
+        """
+        if name is None:
+            name = generate_random_name()
+        return replace(self, name=name)
+
+    def new_sub_action_from(self, task_name: str, input_hash: str, group: str | None) -> ActionID:
+        """Make a deterministic name"""
+        import hashlib
+
+        components = f"{self.run_name}-{self.name}-{input_hash}-{task_name}" + (f"-{group}" if group else "")
+        # has the components into something deterministic
+        bytes_digest = hashlib.md5(components.encode()).digest()
+        new_name = base36_encode(bytes_digest)
+        return self.new_sub_action(new_name)
+
+
+@dataclass(frozen=True, kw_only=True)
+class RawDataPath:
+    """
+    A class representing the raw data path for a task. This is used to store the raw data for the task execution and
+    also get mutations on the path.
+    """
+
+    path: str
+
+    @classmethod
+    def from_local_folder(cls, local_folder: str | pathlib.Path | None = None) -> RawDataPath:
+        """
+        Create a new context attribute object, with local path given. Will be created if it doesn't exist.
+        :return: Path to the temporary directory
+        """
+        match local_folder:
+            case pathlib.Path():
+                local_folder.mkdir(parents=True, exist_ok=True)
+                return RawDataPath(path=str(local_folder))
+            case None:
+                # Create a temporary directory for data storage
+                p = tempfile.mkdtemp()
+                logger.debug(f"Creating temporary directory for data storage: {p}")
+                return RawDataPath(path=p)
+            case str():
+                return RawDataPath(path=local_folder)
+            case _:
+                raise ValueError(f"Invalid local path {local_folder}")
+
+    def get_random_remote_path(self, file_name: Optional[str] = None) -> str:
+        """
+        Returns a random path for uploading a file/directory to.
+
+        :param file_name: If given, will be joined after a randomly generated portion.
+        :return:
+        """
+        import random
+        from uuid import UUID
+
+        import fsspec
+        from fsspec.utils import get_protocol
+
+        random_string = UUID(int=random.getrandbits(128)).hex
+        file_prefix = self.path
+
+        protocol = get_protocol(file_prefix)
+        if "file" in protocol:
+            local_path = pathlib.Path(file_prefix) / random_string
+            if file_name:
+                # Only if file name is given do we create the parent, because it may be needed as a folder otherwise
+                local_path = local_path / file_name
+                if not local_path.exists():
+                    local_path.parent.mkdir(exist_ok=True, parents=True)
+                    local_path.touch()
+            return str(local_path.absolute())
+
+        fs = fsspec.filesystem(protocol)
+        if file_prefix.endswith(fs.sep):
+            file_prefix = file_prefix[:-1]
+        remote_path = fs.sep.join([file_prefix, random_string])
+        if file_name:
+            remote_path = fs.sep.join([remote_path, file_name])
+        return remote_path
+
+
+@dataclass(frozen=True)
+class GroupData:
+    name: str
+
+
+@dataclass(frozen=True, kw_only=True)
+class TaskContext:
+    """
+    A context class to hold the current task executions context.
+    This can be used to access various contextual parameters in the task execution by the user.
+
+    :param action: The action ID of the current execution. This is always set, within a run.
+    :param version: The version of the executed task. This is set when the task is executed by an action and will be
+      set on all sub-actions.
+    """
+
+    action: ActionID
+    version: str
+    raw_data_path: RawDataPath
+    output_path: str
+    run_base_dir: str
+    report: Report
+    group_data: GroupData | None = None
+    checkpoints: Checkpoints | None = None
+    code_bundle: CodeBundle | None = None
+    compiled_image_cache: ImageCache | None = None
+    data: Dict[str, Any] = field(default_factory=dict)
+
+    def replace(self, **kwargs) -> TaskContext:
+        if "data" in kwargs:
+            rec_data = kwargs.pop("data")
+            if rec_data is None:
+                return replace(self, **kwargs)
+            data = {}
+            if self.data is not None:
+                data = self.data.copy()
+            data.update(rec_data)
+            kwargs.update({"data": data})
+        return replace(self, **kwargs)
+
+    def __getitem__(self, key: str) -> Optional[Any]:
+        return self.data.get(key)
+
+
+@dataclass(frozen=True, kw_only=True)
+class CodeBundle:
+    """
+    A class representing a code bundle for a task. This is used to package the code and the inflation path.
+    The code bundle computes the version of the code using the hash of the code.
+
+    :param computed_version: The version of the code bundle. This is the hash of the code.
+    :param destination: The destination path for the code bundle to be inflated to.
+    :param tgz: Optional path to the tgz file.
+    :param pkl: Optional path to the pkl file.
+    :param downloaded_path: The path to the downloaded code bundle. This is only available during runtime, when
+        the code bundle has been downloaded and inflated.
+    """
+
+    computed_version: str
+    destination: str = "."
+    tgz: str | None = None
+    pkl: str | None = None
+    downloaded_path: pathlib.Path | None = None
+
+    # runtime_dependencies: Tuple[str, ...] = field(default_factory=tuple)  In the future if we want we could add this
+    # but this messes up actors, spark etc
+
+    def __post_init__(self):
+        if self.tgz is None and self.pkl is None:
+            raise ValueError("Either tgz or pkl must be provided")
+
+    def with_downloaded_path(self, path: pathlib.Path) -> CodeBundle:
+        """
+        Create a new CodeBundle with the given downloaded path.
+        """
+        return replace(self, downloaded_path=path)
+
+
+@dataclass(frozen=True)
+class Checkpoints:
+    """
+    A class representing the checkpoints for a task. This is used to store the checkpoints for the task execution.
+    """
+
+    prev_checkpoint_path: str | None
+    checkpoint_path: str | None
+
+
+@dataclass(frozen=True)
+class NativeInterface:
+    """
+    A class representing the native interface for a task. This is used to interact with the task and its execution
+    context.
+    """
+
+    inputs: Dict[str, Tuple[Type, Any]]
+    outputs: Dict[str, Type]
+    docstring: Optional[Docstring] = field(default=None)
+
+    def has_outputs(self) -> bool:
+        """
+        Check if the task has outputs. This is used to determine if the task has outputs or not.
+        """
+        return self.outputs is not None and len(self.outputs) > 0
+
+    @classmethod
+    def from_types(cls, inputs: Dict[str, Type], outputs: Dict[str, Type]) -> NativeInterface:
+        """
+        Create a new NativeInterface from the given types. This is used to create a native interface for the task.
+        """
+        return cls(inputs={k: (v, inspect.Parameter.empty) for k, v in inputs.items()}, outputs=outputs)
+
+    @classmethod
+    def from_callable(cls, func: Callable) -> NativeInterface:
+        """
+        Extract the native interface from the given function. This is used to create a native interface for the task.
+        """
+        sig = inspect.signature(func)
+
+        # Extract parameter details (name, type, default value)
+        param_info = {name: (param.annotation, param.default) for name, param in sig.parameters.items()}
+
+        # Get return type
+        outputs = extract_return_annotation(sig.return_annotation)
+        return cls(inputs=param_info, outputs=outputs)
+
+    def convert_to_kwargs(self, *args, **kwargs) -> Dict[str, Any]:
+        """
+        Convert the given arguments to keyword arguments based on the native interface. This is used to convert the
+        arguments to the correct types for the task execution.
+        """
+        # Convert positional arguments to keyword arguments
+        if len(args) > len(self.inputs):
+            raise ValueError(f"Too many positional arguments provided, inputs {self.inputs.keys()}, args {len(args)}")
+        for arg, input_name in zip(args, self.inputs.keys()):
+            kwargs[input_name] = arg
+        return kwargs
+
+    def get_input_types(self) -> Dict[str, Type]:
+        """
+        Get the input types for the task. This is used to get the types of the inputs for the task execution.
+        """
+        return {k: v[0] for k, v in self.inputs.items()}
+
+    def __repr__(self):
+        """
+        Returns a string representation of the task interface.
+        """
+        i = "("
+        if self.inputs:
+            initial = True
+            for key, tpe in self.inputs.items():
+                if not initial:
+                    i += ", "
+                initial = False
+                tp = tpe[0] if isinstance(tpe[0], str) else tpe[0].__name__
+                i += f"{key}: {tp}"
+                if tpe[1] is not inspect.Parameter.empty:
+                    i += f" = {tpe[1]}"
+        i += ")"
+        if self.outputs:
+            initial = True
+            multi = len(self.outputs) > 1
+            i += " -> "
+            if multi:
+                i += "("
+            for key, tpe in self.outputs.items():
+                if not initial:
+                    i += ", "
+                initial = False
+                tp = tpe.__name__ if isinstance(tpe, type) else tpe
+                i += f"{key}: {tp}"
+            if multi:
+                i += ")"
+        return i + ":"
+
+
+@dataclass
+class SerializationContext:
+    """
+    This object holds serialization time contextual information, that can be used when serializing the task and
+    various parameters of a tasktemplate. This is only available when the task is being serialized and can be
+    during a deployment or runtime.
+
+    :param version: The version of the task
+    :param code_bundle: The code bundle for the task. This is used to package the code and the inflation path.
+    :param input_path: The path to the inputs for the task. This is used to determine where the inputs will be located
+    :param output_path: The path to the outputs for the task. This is used to determine where the outputs will be
+     located
+    """
+
+    version: str
+    project: str | None = None
+    domain: str | None = None
+    org: str | None = None
+    code_bundle: Optional[CodeBundle] = None
+    input_path: str = "{{.input}}"
+    output_path: str = "{{.outputPrefix}}"
+    _entrypoint_path: str = field(default="_bin/runtime.py", init=False)
+    image_cache: ImageCache | None = None
+    root_dir: Optional[pathlib.Path] = None
+
+    def get_entrypoint_path(self, interpreter_path: str) -> str:
+        """
+        Get the entrypoint path for the task. This is used to determine the entrypoint for the task execution.
+        :param interpreter_path: The path to the interpreter (python)
+        """
+        return os.path.join(os.path.dirname(os.path.dirname(interpreter_path)), self._entrypoint_path)