flyte 0.0.1b3__py3-none-any.whl → 0.2.0a0__py3-none-any.whl

union/extras/_container.py

@@ -1,263 +0,0 @@
-import os
-import pathlib
-from typing import Any, Dict, List, Literal, Optional, Tuple, Type, Union
-
-from flyteidl.core import tasks_pb2
-
-from union import Image, storage
-from union._datastructures import NativeInterface, SerializationContext
-from union._logging import logger
-from union._task import TaskTemplate
-
-_PRIMARY_CONTAINER_NAME_FIELD = "primary_container_name"
-
-
-def _extract_command_key(cmd: str, **kwargs) -> Any:
-    """
-    Extract the key from the command using regex.
-    """
-    import re
-
-    input_regex = r"^\{\{\s*\.inputs\.(.*?)\s*\}\}$"
-    match = re.match(input_regex, cmd)
-    if match:
-        return match.group(1)
-    return None
-
-
-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 = False,
-        **kwargs,
-    ):
-        super().__init__(
-            task_type="raw-container",
-            name=name,
-            image=image,
-            interface=NativeInterface({k: (v, None) for k, v in inputs.items()}, outputs),
-            **kwargs,
-        )
-        self._image = image
-        self._cmd = command
-        self._args = arguments
-        self._input_data_dir = input_data_dir
-        self._output_data_dir = 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 flytekit.types.directory import FlyteDirectory
-        # from flytekit.types.file import FlyteFile
-
-        volume_binding = {}
-        path_k = _extract_path_command_key(cmd, self._input_data_dir)
-        k = path_k if path_k else _extract_command_key(cmd)
-
-        if k:
-            input_val = kwargs.get(k)
-            # TODO: Add support file and directory transformer first
-            # if type(input_val) in [FlyteFile, FlyteDirectory]:
-            #     if not path_k:
-            #         raise AssertionError(
-            #             "FlyteFile and FlyteDirectory 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
-            command = str(input_val)
-        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: str) -> Dict[str, Any]:
-        output_dict = {}
-        if self._outputs:
-            for k, output_type in self._outputs.items():
-                output_path = os.path.join(output_directory, k)
-                with open(output_path, "r") as f:
-                    output_val = f.read()
-                output_dict[k] = self._convert_output_val_to_correct_type(output_val, output_type)
-        return output_dict
-
-    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[output_directory] = {"bind": self._output_data_dir, "mode": "rw"}
-
-        client = docker.from_env()
-        self._pull_image_if_not_exists(client, self._image)
-
-        container = client.containers.run(
-            self._image, 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()}")
-
-        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=self._input_data_dir,
-            output_path=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 [])
-
-    def config(self, sctx: SerializationContext) -> Optional[Dict[str, str]]:
-        if self.pod_template is None:
-            return {}
-        return {_PRIMARY_CONTAINER_NAME_FIELD: self.primary_container_name}

union/io/__init__.py

@@ -1,11 +0,0 @@
-"""
-## 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"]
-
-from ._dir import Dir
-from ._file import File

union/io/_dir.py

@@ -1,425 +0,0 @@
-from __future__ import annotations
-
-import os
-from pathlib import Path
-from typing import AsyncIterator, Generic, Iterator, List, Optional, Type, TypeVar, Union
-
-from flyteidl.core import literals_pb2, types_pb2
-from fsspec.asyn import AsyncFileSystem
-from pydantic import BaseModel, model_validator
-
-import union.storage as storage
-from union.io._file import File
-from union.types import TypeEngine, TypeTransformer, TypeTransformerFailedError
-
-# Type variable for the directory format
-T = TypeVar("T")
-
-
-class Dir(BaseModel, Generic[T]):
-    """
-    A generic directory class representing a directory with files of a specified format.
-    Provides both async and sync interfaces for directory operations.
-    Users are responsible for handling all I/O - the type transformer for Dir does not do any automatic uploading
-    or downloading of files.
-
-    The generic type T represents the format of the files in the directory.
-
-    Example:
-        ```python
-        # Async usage
-        from pandas import DataFrame
-        data_dir = Dir[DataFrame](path="s3://my-bucket/data/")
-
-        # Walk through files
-        async for file in data_dir.walk():
-            async with file.open() as f:
-                content = await f.read()
-
-        # Sync alternative
-        for file in data_dir.walk_sync():
-            with file.open_sync() as f:
-                content = f.read()
-        ```
-    """
-
-    # Represents either a local or remote path.
-    path: str
-    name: Optional[str] = None
-    format: str = ""
-
-    class Config:
-        arbitrary_types_allowed = True
-
-    @model_validator(mode="before")
-    @classmethod
-    def pre_init(cls, data):
-        if data.get("name") is None:
-            data["name"] = Path(data["path"]).name
-        return data
-
-    async def walk(self, recursive: bool = True, max_depth: Optional[int] = None) -> AsyncIterator[File[T]]:
-        """
-        Asynchronously walk through the directory and yield File objects.
-
-        Args:
-            recursive: If True, recursively walk subdirectories
-            max_depth: Maximum depth for recursive walking
-
-        Yields:
-            File objects for each file found in the directory
-
-        Example:
-            ```python
-            async for file in directory.walk():
-                local_path = await file.download()
-                # Process the file
-            ```
-        """
-        fs = storage.get_underlying_filesystem(path=self.path)
-        if recursive is False:
-            max_depth = 2
-
-        # Note if the path is actually just a file, no walking is done.
-        if isinstance(fs, AsyncFileSystem):
-            async for parent, _, files in fs._walk(self.path, maxdepth=max_depth):
-                for file in files:
-                    full_file = fs.unstrip_protocol(parent + fs.sep + file)
-                    yield File[T](path=full_file)
-        else:
-            for parent, _, files in fs.walk(self.path, maxdepth=max_depth):
-                for file in files:
-                    if "file" in fs.protocol:
-                        full_file = os.path.join(parent, file)
-                    else:
-                        full_file = fs.unstrip_protocol(parent + fs.sep + file)
-                    yield File[T](path=full_file)
-
-    def walk_sync(
-        self, recursive: bool = True, file_pattern: str = "*", max_depth: Optional[int] = None
-    ) -> Iterator[File[T]]:
-        """
-        Synchronously walk through the directory and yield File objects.
-
-        Args:
-            recursive: If True, recursively walk subdirectories
-            file_pattern: Glob pattern to filter files
-            max_depth: Maximum depth for recursive walking
-
-        Yields:
-            File objects for each file found in the directory
-
-        Example:
-            ```python
-            for file in directory.walk_sync():
-                local_path = file.download_sync()
-                # Process the file
-            ```
-        """
-        fs = storage.get_underlying_filesystem(path=self.path)
-        for parent, _, files in fs.walk(self.path, maxdepth=max_depth):
-            for file in files:
-                if "file" in fs.protocol:
-                    full_file = os.path.join(parent, file)
-                else:
-                    full_file = fs.unstrip_protocol(parent + fs.sep + file)
-                yield File[T](path=full_file)
-
-    async def list_files(self) -> List[File[T]]:
-        """
-        Asynchronously get a list of all files in the directory (non-recursive).
-
-        Returns:
-            A list of File objects
-
-        Example:
-            ```python
-            files = await directory.list_files()
-            for file in files:
-                # Process the file
-            ```
-        """
-        # todo: this should probably also just defer to fsspec.find()
-        files = []
-        async for file in self.walk(recursive=False):
-            files.append(file)
-        return files
-
-    def list_files_sync(self) -> List[File[T]]:
-        """
-        Synchronously get a list of all files in the directory (non-recursive).
-
-        Returns:
-            A list of File objects
-
-        Example:
-            ```python
-            files = directory.list_files_sync()
-            for file in files:
-                # Process the file
-            ```
-        """
-        return list(self.walk_sync(recursive=False))
-
-    async def download(self, local_path: Optional[Union[str, Path]] = None) -> str:
-        """
-        Asynchronously download the entire directory to a local path.
-
-        Args:
-            local_path: The local path to download the directory to. If None, a temporary
-                       directory will be used.
-
-        Returns:
-            The path to the downloaded directory
-
-        Example:
-            ```python
-            local_dir = await directory.download('/tmp/my_data/')
-            ```
-        """
-        local_dest = str(local_path) if local_path else str(storage.get_random_local_path())
-        if not storage.is_remote(self.path):
-            if not local_path or local_path == self.path:
-                # Skip copying
-                return self.path
-            else:
-                # Shell out to a thread to copy
-                import asyncio
-                import shutil
-
-                async def copy_tree():
-                    loop = asyncio.get_event_loop()
-                    await loop.run_in_executor(None, lambda: shutil.copytree(self.path, local_dest, dirs_exist_ok=True))
-
-                await copy_tree()
-        return await storage.get(self.path, local_dest, recursive=True)
-
-    def download_sync(self, local_path: Optional[Union[str, Path]] = None) -> str:
-        """
-        Synchronously download the entire directory to a local path.
-
-        Args:
-            local_path: The local path to download the directory to. If None, a temporary
-                       directory will be used.
-
-        Returns:
-            The path to the downloaded directory
-
-        Example:
-            ```python
-            local_dir = directory.download_sync('/tmp/my_data/')
-            ```
-        """
-        local_dest = str(local_path) if local_path else str(storage.get_random_local_path())
-        if not storage.is_remote(self.path):
-            if not local_path or local_path == self.path:
-                # Skip copying
-                return self.path
-            else:
-                # Shell out to a thread to copy
-                import shutil
-
-                shutil.copytree(self.path, local_dest, dirs_exist_ok=True)
-
-        # Figure this out when we figure out the final synchronicity story
-        raise NotImplementedError("Sync download is not implemented for remote paths")
-
-    @classmethod
-    async def from_local(cls, local_path: Union[str, Path], remote_path: Optional[str] = None) -> Dir[T]:
-        """
-        Asynchronously create a new Dir by uploading a local directory to the configured remote store.
-
-        Args:
-            local_path: Path to the local directory
-            remote_path: Optional path to store the directory remotely. If None, a path will be generated.
-
-        Returns:
-            A new Dir instance pointing to the uploaded directory
-
-        Example:
-            ```python
-            remote_dir = await Dir[DataFrame].from_local('/tmp/data_dir/', 's3://bucket/data/')
-            ```
-        """
-        local_path_str = str(local_path)
-        dirname = os.path.basename(os.path.normpath(local_path_str))
-
-        output_path = await storage.put(from_path=local_path_str, to_path=remote_path, recursive=True)
-        return cls(path=output_path, name=dirname)
-
-    @classmethod
-    def from_local_sync(cls, local_path: Union[str, Path], remote_path: Optional[str] = None) -> Dir[T]:
-        """
-        Synchronously create a new Dir by uploading a local directory to the configured remote store.
-
-        Args:
-            local_path: Path to the local directory
-            remote_path: Optional path to store the directory remotely. If None, a path will be generated.
-
-        Returns:
-            A new Dir instance pointing to the uploaded directory
-
-        Example:
-            ```python
-            remote_dir = Dir[DataFrame].from_local_sync('/tmp/data_dir/', 's3://bucket/data/')
-            ```
-        """
-        # Implement this after we figure out the final synchronicity story
-        raise NotImplementedError("Sync upload is not implemented for remote paths")
-
-    async def exists(self) -> bool:
-        """
-        Asynchronously check if the directory exists.
-
-        Returns:
-            True if the directory exists, False otherwise
-
-        Example:
-            ```python
-            if await directory.exists():
-                # Process the directory
-            ```
-        """
-        fs = storage.get_underlying_filesystem(path=self.path)
-        if isinstance(fs, AsyncFileSystem):
-            return await fs._exists(self.path)
-        else:
-            return fs.exists(self.path)
-
-    def exists_sync(self) -> bool:
-        """
-        Synchronously check if the directory exists.
-
-        Returns:
-            True if the directory exists, False otherwise
-
-        Example:
-            ```python
-            if directory.exists_sync():
-                # Process the directory
-            ```
-        """
-        fs = storage.get_underlying_filesystem(path=self.path)
-        return fs.exists(self.path)
-
-    async def get_file(self, file_name: str) -> Optional[File[T]]:
-        """
-        Asynchronously get a specific file from the directory.
-
-        Args:
-            file_name: The name of the file to get
-
-        Returns:
-            A File instance if the file exists, None otherwise
-
-        Example:
-            ```python
-            file = await directory.get_file("data.csv")
-            if file:
-                # Process the file
-            ```
-        """
-        fs = storage.get_underlying_filesystem(path=self.path)
-        file_path = fs.sep.join([self.path, file_name])
-        file = File[T](path=file_path)
-
-        if fs.exists(file_path):
-            return file
-        return None
-
-    def get_file_sync(self, file_name: str) -> Optional[File[T]]:
-        """
-        Synchronously get a specific file from the directory.
-
-        Args:
-            file_name: The name of the file to get
-
-        Returns:
-            A File instance if the file exists, None otherwise
-
-        Example:
-            ```python
-            file = directory.get_file_sync("data.csv")
-            if file:
-                # Process the file
-            ```
-        """
-        file_path = os.path.join(self.path, file_name)
-        file = File[T](path=file_path)
-
-        if file.exists_sync():
-            return file
-        return None
-
-
-class DirTransformer(TypeTransformer[Dir]):
-    """
-    Transformer for Dir objects. This type transformer does not handle any i/o. That is now the responsibility of the
-    user.
-    """
-
-    def __init__(self):
-        super().__init__(name="Dir", t=Dir)
-
-    def get_literal_type(self, t: Type[Dir]) -> types_pb2.LiteralType:
-        """Get the Flyte literal type for a File type."""
-        return types_pb2.LiteralType(
-            blob=types_pb2.BlobType(
-                # todo: set format from generic
-                format="",  # Format is determined by the generic type T
-                dimensionality=types_pb2.BlobType.BlobDimensionality.MULTIPART,
-            )
-        )
-
-    async def to_literal(
-        self,
-        python_val: Dir,
-        python_type: Type[Dir],
-        expected: types_pb2.LiteralType,
-    ) -> literals_pb2.Literal:
-        """Convert a Dir object to a Flyte literal."""
-        if not isinstance(python_val, Dir):
-            raise TypeTransformerFailedError(f"Expected Dir object, received {type(python_val)}")
-
-        return literals_pb2.Literal(
-            scalar=literals_pb2.Scalar(
-                blob=literals_pb2.Blob(
-                    metadata=literals_pb2.BlobMetadata(
-                        type=types_pb2.BlobType(
-                            format=python_val.format, dimensionality=types_pb2.BlobType.BlobDimensionality.MULTIPART
-                        )
-                    ),
-                    uri=python_val.path,
-                )
-            )
-        )
-
-    async def to_python_value(
-        self,
-        lv: literals_pb2.Literal,
-        expected_python_type: Type[Dir],
-    ) -> Dir:
-        """Convert a Flyte literal to a File object."""
-        if not lv.scalar.HasField("blob"):
-            raise TypeTransformerFailedError(f"Expected blob literal, received {lv}")
-        if not lv.scalar.blob.metadata.type.dimensionality == types_pb2.BlobType.BlobDimensionality.MULTIPART:
-            raise TypeTransformerFailedError(
-                f"Expected multipart, received {lv.scalar.blob.metadata.type.dimensionality}"
-            )
-
-        uri = lv.scalar.blob.uri
-        filename = Path(uri).name
-        f = Dir(path=uri, name=filename, format=lv.scalar.blob.metadata.type.format)
-        return f
-
-    def guess_python_type(self, literal_type: types_pb2.LiteralType) -> Type[Dir]:
-        """Guess the Python type from a Flyte literal type."""
-        if (
-            literal_type.HasField("blob")
-            and literal_type.blob.dimensionality == types_pb2.BlobType.BlobDimensionality.MULTIPART
-        ):
-            return Dir
-        raise ValueError(f"Cannot guess python type from {literal_type}")
-
-
-TypeEngine.register(DirTransformer())