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

flyte/config/_reader.py

@@ -0,0 +1,207 @@
+import os
+import pathlib
+import typing
+from dataclasses import dataclass
+from functools import lru_cache
+from os import getenv
+from pathlib import Path
+
+import yaml
+
+from flyte._logging import logger
+
+# This is the default config file name for flyte
+FLYTECTL_CONFIG_ENV_VAR = "FLYTECTL_CONFIG"
+UCTL_CONFIG_ENV_VAR = "UCTL_CONFIG"
+
+
+@dataclass
+class YamlConfigEntry(object):
+    """
+    Creates a record for the config entry.
+    Args:
+        switch: dot-delimited string that should match flytectl args. Leaving it as dot-delimited instead of a list
+          of strings because it's easier to maintain alignment with flytectl.
+        config_value_type: Expected type of the value
+    """
+
+    switch: str
+    config_value_type: typing.Type = str
+
+    def get_env_name(self) -> str:
+        var_name = self.switch.upper().replace(".", "_")
+        return f"FLYTE_{var_name}"
+
+    def read_from_env(self, transform: typing.Optional[typing.Callable] = None) -> typing.Optional[typing.Any]:
+        """
+        Reads the config entry from environment variable, the structure of the env var is current
+        ``FLYTE_{SECTION}_{OPTION}`` all upper cased. We will change this in the future.
+        :return:
+        """
+        env = self.get_env_name()
+        v = os.environ.get(env, None)
+        if v is None:
+            return None
+        return transform(v) if transform else v
+
+    def read_from_file(
+        self, cfg: "ConfigFile", transform: typing.Optional[typing.Callable] = None
+    ) -> typing.Optional[typing.Any]:
+        if not cfg:
+            return None
+        try:
+            v = cfg.get(self)
+            if isinstance(v, bool) or bool(v is not None and v):
+                return transform(v) if transform else v
+        except Exception:
+            ...
+
+        return None
+
+
+@dataclass
+class ConfigEntry(object):
+    """
+    A top level Config entry holder, that holds multiple different representations of the config.
+    Legacy means the INI style config files. YAML support is for the flytectl config file, which is there by default
+    when flytectl starts a sandbox
+    """
+
+    yaml_entry: YamlConfigEntry
+    transform: typing.Optional[typing.Callable[[str], typing.Any]] = None
+
+    def read(self, cfg: typing.Optional["ConfigFile"] = None) -> typing.Optional[typing.Any]:
+        """
+        Reads the config Entry from the various sources in the following order,
+        #. First try to read from the relevant environment variable,
+        #. If missing, then try to read from the legacy config file, if one was parsed.
+        #. If missing, then try to read from the yaml file.
+
+        The constructor for ConfigFile currently does not allow specification of both the ini and yaml style formats.
+
+        :param cfg:
+        :return:
+        """
+        from_env = self.yaml_entry.read_from_env(self.transform)
+        if from_env is not None:
+            return from_env
+        if cfg and cfg.yaml_config and self.yaml_entry:
+            return self.yaml_entry.read_from_file(cfg, self.transform)
+
+        return None
+
+
+class ConfigFile(object):
+    def __init__(self, location: str):
+        """
+        Load the config from this location
+        """
+        self._location = location
+        self._yaml_config = self._read_yaml_config(location)
+
+    @property
+    def path(self) -> pathlib.Path:
+        """
+        Returns the path to the config file.
+        :return: Path to the config file
+        """
+        return pathlib.Path(self._location)
+
+    @staticmethod
+    def _read_yaml_config(location: str) -> typing.Optional[typing.Dict[str, typing.Any]]:
+        with open(location, "r") as fh:
+            try:
+                yaml_contents = yaml.safe_load(fh)
+                return yaml_contents
+            except yaml.YAMLError as exc:
+                logger.warning(f"Error {exc} reading yaml config file at {location}, ignoring...")
+                return None
+
+    def _get_from_yaml(self, c: YamlConfigEntry) -> typing.Any:
+        keys = c.switch.split(".")  # flytectl switches are dot delimited
+        d = typing.cast(typing.Dict[str, typing.Any], self.yaml_config)
+        try:
+            for k in keys:
+                d = d[k]
+            return d
+        except KeyError:
+            return None
+
+    def get(self, c: YamlConfigEntry) -> typing.Any:
+        return self._get_from_yaml(c)
+
+    @property
+    def yaml_config(self) -> typing.Dict[str, typing.Any] | None:
+        return self._yaml_config
+
+
+def resolve_config_path() -> pathlib.Path | None:
+    """
+    Config is read from the following locations in order of precedence:
+    1. ./config.yaml if it exists
+    2. `UCTL_CONFIG` environment variable
+    3. `FLYTECTL_CONFIG` environment variable
+    4. ~/.union/config.yaml if it exists
+    5. ~/.flyte/config.yaml if it exists
+    """
+    current_location_config = Path("config.yaml")
+    if current_location_config.exists():
+        return current_location_config
+    logger.debug("No ./config.yaml found, returning None")
+
+    uctl_path_from_env = getenv(UCTL_CONFIG_ENV_VAR, None)
+    if uctl_path_from_env:
+        return pathlib.Path(uctl_path_from_env)
+    logger.debug("No UCTL_CONFIG environment variable found, checking FLYTECTL_CONFIG")
+
+    flytectl_path_from_env = getenv(FLYTECTL_CONFIG_ENV_VAR, None)
+    if flytectl_path_from_env:
+        return pathlib.Path(flytectl_path_from_env)
+    logger.debug("No FLYTECTL_CONFIG environment variable found, checking default locations")
+
+    home_dir_union_config = Path(Path.home(), ".union", "config.yaml")
+    if home_dir_union_config.exists():
+        return home_dir_union_config
+    logger.debug("No ~/.union/config.yaml found, checking current directory")
+
+    home_dir_flytectl_config = Path(Path.home(), ".flyte", "config.yaml")
+    if home_dir_flytectl_config.exists():
+        return home_dir_flytectl_config
+    logger.debug("No ~/.flyte/config.yaml found, checking current directory")
+
+    return None
+
+
+@lru_cache
+def get_config_file(c: typing.Union[str, ConfigFile, None]) -> ConfigFile | None:
+    """
+    Checks if the given argument is a file or a configFile and returns a loaded configFile else returns None
+    """
+    if "PYTEST_VERSION" in os.environ:
+        # Use default local config in the pytest environment
+        return None
+    if isinstance(c, str):
+        logger.debug(f"Using specified config file at {c}")
+        return ConfigFile(c)
+    elif isinstance(c, ConfigFile):
+        return c
+    config_path = resolve_config_path()
+    if config_path:
+        return ConfigFile(str(config_path))
+    return None
+
+
+def read_file_if_exists(filename: typing.Optional[str], encoding=None) -> typing.Optional[str]:
+    """
+    Reads the contents of the file if passed a path. Otherwise, returns None.
+
+    :param filename: The file path to load
+    :param encoding: The encoding to use when reading the file.
+    :return: The contents of the file as a string or None.
+    """
+    if not filename:
+        return None
+
+    file = pathlib.Path(filename)
+    logger.debug(f"Reading file contents from [{file}] with current directory [{os.getcwd()}].")
+    return file.read_text(encoding=encoding)

flyte/errors.py

@@ -0,0 +1,172 @@
+"""
+Exceptions raised by Union.
+
+These errors are raised when the underlying task execution fails, either because of a user error, system error or an
+unknown error.
+"""
+
+from typing import Literal
+
+ErrorKind = Literal["system", "unknown", "user"]
+
+
+class BaseRuntimeError(RuntimeError):
+    """
+    Base class for all Union runtime errors. These errors are raised when the underlying task execution fails, either
+    because of a user error, system error or an unknown error.
+    """
+
+    def __init__(self, code: str, kind: ErrorKind, root_cause_message: str, worker: str | None = None):
+        super().__init__(root_cause_message)
+        self.code = code
+        self.kind = kind
+        self.worker = worker
+
+
+class InitializationError(BaseRuntimeError):
+    """
+    This error is raised when the Union system is tried to access without being initialized.
+    """
+
+
+class RuntimeSystemError(BaseRuntimeError):
+    """
+    This error is raised when the underlying task execution fails because of a system error. This could be a bug in the
+    Union system or a bug in the user's code.
+    """
+
+    def __init__(self, code: str, message: str, worker: str | None = None):
+        super().__init__(code, "system", message, worker)
+
+
+class UnionRpcError(RuntimeSystemError):
+    """
+    This error is raised when communication with the Union server fails.
+    """
+
+
+class RuntimeUserError(BaseRuntimeError):
+    """
+    This error is raised when the underlying task execution fails because of an error in the user's code.
+    """
+
+    def __init__(self, code: str, message: str, worker: str | None = None):
+        super().__init__(code, "user", message, worker)
+
+
+class RuntimeUnknownError(BaseRuntimeError):
+    """
+    This error is raised when the underlying task execution fails because of an unknown error.
+    """
+
+    def __init__(self, code: str, message: str, worker: str | None = None):
+        super().__init__(code, "unknown", message, worker)
+
+
+class OOMError(RuntimeUserError):
+    """
+    This error is raised when the underlying task execution fails because of an out-of-memory error.
+    """
+
+
+class TaskInterruptedError(RuntimeUserError):
+    """
+    This error is raised when the underlying task execution is interrupted.
+    """
+
+
+class PrimaryContainerNotFoundError(RuntimeUserError):
+    """
+    This error is raised when the primary container is not found.
+    """
+
+
+class TaskTimeoutError(RuntimeUserError):
+    """
+    This error is raised when the underlying task execution runs for longer than the specified timeout.
+    """
+
+
+class RetriesExhaustedError(RuntimeUserError):
+    """
+    This error is raised when the underlying task execution fails after all retries have been exhausted.
+    """
+
+
+class InvalidImageNameError(RuntimeUserError):
+    """
+    This error is raised when the image name is invalid.
+    """
+
+
+class ImagePullBackOffError(RuntimeUserError):
+    """
+    This error is raised when the image cannot be pulled.
+    """
+
+
+class CustomError(RuntimeUserError):
+    """
+    This error is raised when the user raises a custom error.
+    """
+
+    def __init__(self, code: str, message: str):
+        super().__init__(code, message, "user")
+
+    @classmethod
+    def from_exception(cls, e: Exception):
+        """
+        Create a CustomError from an exception. The exception's class name is used as the error code and the exception
+        message is used as the error message.
+        """
+        return cls(e.__class__.__name__, str(e))
+
+
+class NotInTaskContextError(RuntimeUserError):
+    """
+    This error is raised when the user tries to access the task context outside of a task.
+    """
+
+
+class ActionNotFoundError(RuntimeError):
+    """
+    This error is raised when the user tries to access an action that does not exist.
+    """
+
+
+class ReferenceTaskError(RuntimeUserError):
+    """
+    This error is raised when the user tries to access a task that does not exist.
+    """
+
+    def __init__(self, message: str):
+        super().__init__("ReferenceTaskUsageError", message, "user")
+
+
+class LogsNotYetAvailableError(BaseRuntimeError):
+    """
+    This error is raised when the logs are not yet available for a task.
+    """
+
+    def __init__(self, message: str):
+        super().__init__("LogsNotYetAvailable", "system", message, None)
+
+
+class RuntimeDataValidationError(RuntimeUserError):
+    """
+    This error is raised when the user tries to access a resource that does not exist or is invalid.
+    """
+
+    def __init__(self, var: str, e: Exception, task_name: str = ""):
+        super().__init__(
+            "DataValiationError", f"In task {task_name} variable {var}, failed to serialize/deserialize because {e}"
+        )
+
+
+class DeploymentError(RuntimeUserError):
+    """
+    This error is raised when the deployment of a task fails, or some preconditions for deployment are not met.
+    """
+
+    def __init__(self, message: str):
+        super().__init__("DeploymentError", message, "user")

flyte/extras/__init__.py

@@ -0,0 +1,5 @@
+from ._container import ContainerTask
+
+__all__ = [
+    "ContainerTask",
+]

flyte/extras/_container.py

@@ -0,0 +1,263 @@
+import os
+import pathlib
+from typing import Any, Dict, List, Literal, Optional, Tuple, Type, Union
+
+from flyteidl.core import tasks_pb2
+
+from flyte import Image, storage
+from flyte._logging import logger
+from flyte._task import TaskTemplate
+from flyte.models import NativeInterface, SerializationContext
+
+
+def _extract_command_key(cmd: str, **kwargs) -> List[Any] | None:
+    """
+    Extract the key from the command using regex.
+    """
+    import re
+
+    input_regex = r"\{\{\.inputs\.([a-zA-Z0-9_]+)\}\}"
+    return re.findall(input_regex, cmd)
+
+
+def _extract_path_command_key(cmd: str, input_data_dir: Optional[str]) -> Optional[str]:
+    """
+    Extract the key from the path-like command using regex.
+    """
+    import re
+
+    input_data_dir = input_data_dir or ""
+    input_regex = rf"{re.escape(input_data_dir)}/(.+)$"
+    match = re.match(input_regex, cmd)
+    if match:
+        return match.group(1)
+    return None
+
+
+class ContainerTask(TaskTemplate):
+    """
+    This is an intermediate class that represents Flyte Tasks that run a container at execution time. This is the vast
+    majority of tasks - the typical ``@task`` decorated tasks; for instance, all run a container. An example of
+    something that doesn't run a container would be something like the Athena SQL task.
+
+    :param name: Name of the task
+    :param image: The container image to use for the task. This can be a string or an Image object.
+    :param command: The command to run in the container. This can be a list of strings or a single string.
+    :param inputs: The inputs to the task. This is a dictionary of input names to types.
+    :param arguments: The arguments to pass to the command. This is a list of strings.
+    :param outputs: The outputs of the task. This is a dictionary of output names to types.
+    :param input_data_dir: The directory where the input data is stored. This is a string or a Path object.
+    :param output_data_dir: The directory where the output data is stored. This is a string or a Path object.
+    :param metadata_format: The format of the output file. This can be "JSON", "YAML", or "PROTO".
+    :param local_logs: If True, logs will be printed to the console in the local execution.
+    """
+
+    MetadataFormat = Literal["JSON", "YAML", "PROTO"]
+
+    def __init__(
+        self,
+        name: str,
+        image: Union[str, Image],
+        command: List[str],
+        inputs: Optional[Dict[str, Type]] = None,
+        arguments: Optional[List[str]] = None,
+        outputs: Optional[Dict[str, Type]] = None,
+        input_data_dir: str | pathlib.Path = "/var/inputs",
+        output_data_dir: str | pathlib.Path = "/var/outputs",
+        metadata_format: MetadataFormat = "JSON",
+        local_logs: bool = True,
+        **kwargs,
+    ):
+        super().__init__(
+            task_type="raw-container",
+            name=name,
+            image=image,
+            interface=NativeInterface({k: (v, None) for k, v in inputs.items()} if inputs else {}, outputs or {}),
+            **kwargs,
+        )
+        self._image = image
+        if isinstance(image, str):
+            if image == "auto":
+                self._image = Image.auto()
+            else:
+                self._image = Image.from_prebuilt(image)
+        self._cmd = command
+        self._args = arguments
+        self._input_data_dir = input_data_dir
+        if isinstance(input_data_dir, str):
+            self._input_data_dir = pathlib.Path(input_data_dir)
+        self._output_data_dir = output_data_dir
+        if isinstance(output_data_dir, str):
+            self._output_data_dir = pathlib.Path(output_data_dir)
+        self._metadata_format = metadata_format
+        self._inputs = inputs
+        self._outputs = outputs
+        self.local_logs = local_logs
+
+    def _render_command_and_volume_binding(self, cmd: str, **kwargs) -> Tuple[str, Dict[str, Dict[str, str]]]:
+        """
+        We support template-style references to inputs, e.g., "{{.inputs.infile}}".
+
+        For FlyteFile and FlyteDirectory commands, e.g., "/var/inputs/inputs", we extract the key from strings that
+         begin with the specified `input_data_dir`.
+        """
+        from flyte.io import Dir, File
+
+        volume_binding: Dict[str, Dict[str, str]] = {}
+        path_k = _extract_path_command_key(cmd, str(self._input_data_dir))
+        keys = path_k if path_k else _extract_command_key(cmd)
+
+        if keys:
+            for k in keys:
+                input_val = kwargs.get(k)
+                # TODO: Add support file and directory transformer first
+                if type(input_val) in [File, Dir]:
+                    if not path_k:
+                        raise AssertionError(
+                            "File and Directory commands should not use the template syntax "
+                            "like this: {{.inputs.infile}}\n"
+                            "Please use a path-like syntax, such as: /var/inputs/infile.\n"
+                            "This requirement is due to how Flyte Propeller processes template syntax inputs."
+                        )
+                    local_flyte_file_or_dir_path = str(input_val)
+                    remote_flyte_file_or_dir_path = os.path.join(self._input_data_dir, k)  # type: ignore
+                    volume_binding[local_flyte_file_or_dir_path] = {
+                        "bind": remote_flyte_file_or_dir_path,
+                        "mode": "rw",
+                    }
+                    command = remote_flyte_file_or_dir_path
+                else:
+                    command = cmd
+
+        return command, volume_binding
+
+    def _prepare_command_and_volumes(
+        self, cmd_and_args: List[str], **kwargs
+    ) -> Tuple[List[str], Dict[str, Dict[str, str]]]:
+        """
+        Prepares the command and volume bindings for the container based on input arguments and command templates.
+
+        Parameters:
+        - cmd_and_args (List[str]): The command and arguments to prepare.
+        - **kwargs: Keyword arguments representing task inputs.
+
+        Returns:
+        - Tuple[List[str], Dict[str, Dict[str, str]]]: A tuple containing the prepared commands and volume bindings.
+        """
+
+        commands = []
+        volume_bindings = {}
+
+        for cmd in cmd_and_args:
+            command, volume_binding = self._render_command_and_volume_binding(cmd, **kwargs)
+            commands.append(command)
+            volume_bindings.update(volume_binding)
+
+        return commands, volume_bindings
+
+    def _pull_image_if_not_exists(self, client, image: str):
+        try:
+            if not client.images.list(filters={"reference": image}):
+                logger.info(f"Pulling image: {image} for container task: {self.name}")
+                client.images.pull(image)
+        except Exception as e:
+            logger.error(f"Failed to pull image {image}: {e!s}")
+            raise
+
+    def _string_to_timedelta(self, s: str):
+        import datetime
+        import re
+
+        regex = r"(?:(\d+) days?, )?(?:(\d+):)?(\d+):(\d+)(?:\.(\d+))?"
+        parts = re.match(regex, s)
+        if not parts:
+            raise ValueError("Invalid timedelta string format")
+
+        days = int(parts.group(1)) if parts.group(1) else 0
+        hours = int(parts.group(2)) if parts.group(2) else 0
+        minutes = int(parts.group(3)) if parts.group(3) else 0
+        seconds = int(parts.group(4)) if parts.group(4) else 0
+        microseconds = int(parts.group(5)) if parts.group(5) else 0
+
+        return datetime.timedelta(
+            days=days,
+            hours=hours,
+            minutes=minutes,
+            seconds=seconds,
+            microseconds=microseconds,
+        )
+
+    def _convert_output_val_to_correct_type(self, output_val: Any, output_type: Type) -> Any:
+        import datetime
+
+        if issubclass(output_type, bool):
+            return output_val.lower() != "false"
+        elif issubclass(output_type, datetime.datetime):
+            return datetime.datetime.fromisoformat(output_val)
+        elif issubclass(output_type, datetime.timedelta):
+            return self._string_to_timedelta(output_val)
+        else:
+            return output_type(output_val)
+
+    def _get_output_dict(self, output_directory: pathlib.Path) -> Dict[str, Any]:
+        output_dict = {}
+        if self._outputs:
+            for k, output_type in self._outputs.items():
+                output_path = output_directory / k
+                with output_path.open("r") as f:
+                    output_val = f.read()
+                output_dict[k] = self._convert_output_val_to_correct_type(output_val, output_type)
+        return output_dict
+
+    async def execute(self, **kwargs) -> Any:
+        try:
+            import docker
+        except ImportError:
+            raise ImportError("Docker is not installed. Please install Docker by running `pip install docker`.")
+
+        # Normalize the input and output directories
+        self._input_data_dir = os.path.normpath(self._input_data_dir) if self._input_data_dir else ""
+        self._output_data_dir = os.path.normpath(self._output_data_dir) if self._output_data_dir else ""
+
+        output_directory = storage.get_random_local_directory()
+        cmd_and_args = (self._cmd or []) + (self._args or [])
+        commands, volume_bindings = self._prepare_command_and_volumes(cmd_and_args, **kwargs)
+        volume_bindings[str(output_directory)] = {"bind": self._output_data_dir, "mode": "rw"}
+
+        client = docker.from_env()
+        if isinstance(self._image, str):
+            raise AssertionError(f"Only Image objects are supported, not strings. Got {self._image} instead.")
+        uri = self._image.uri
+        self._pull_image_if_not_exists(client, uri)
+        print(f"Command: {commands!r}")
+
+        container = client.containers.run(uri, command=commands, remove=True, volumes=volume_bindings, detach=True)
+
+        # Wait for the container to finish the task
+        # TODO: Add a 'timeout' parameter to control the max wait time for the container to finish the task.
+
+        if self.local_logs:
+            for log in container.logs(stream=True):
+                print(f"[Local Container] {log.strip()!r}")
+
+        container.wait()
+
+        output_dict = self._get_output_dict(output_directory)
+        return output_dict
+
+    def data_loading_config(self, sctx: SerializationContext) -> tasks_pb2.DataLoadingConfig:
+        literal_to_protobuf = {
+            "JSON": tasks_pb2.DataLoadingConfig.JSON,
+            "YAML": tasks_pb2.DataLoadingConfig.YAML,
+            "PROTO": tasks_pb2.DataLoadingConfig.PROTO,
+        }
+
+        return tasks_pb2.DataLoadingConfig(
+            input_path=str(self._input_data_dir),
+            output_path=str(self._output_data_dir),
+            enabled=True,
+            format=literal_to_protobuf.get(self._metadata_format, "JSON"),
+        )
+
+    def container_args(self, sctx: SerializationContext) -> List[str]:
+        return self._cmd + (self._args if self._args else [])

flyte/io/__init__.py

@@ -0,0 +1,27 @@
+"""
+## IO data types
+
+This package contains additional data types beyond the primitive data types in python to abstract data flow
+of large datasets in Union.
+
+"""
+
+__all__ = [
+    "Dir",
+    "File",
+    "StructuredDataset",
+    "StructuredDatasetDecoder",
+    "StructuredDatasetEncoder",
+    "StructuredDatasetTransformerEngine",
+    "lazy_import_structured_dataset_handler",
+]
+
+from ._dir import Dir
+from ._file import File
+from ._structured_dataset import (
+    StructuredDataset,
+    StructuredDatasetDecoder,
+    StructuredDatasetEncoder,
+    StructuredDatasetTransformerEngine,
+    lazy_import_structured_dataset_handler,
+)