flyte 0.0.1b0__py3-none-any.whl

union/storage/_storage.py

@@ -0,0 +1,247 @@
+import os
+import pathlib
+import random
+import tempfile
+import typing
+from typing import AsyncIterator, Optional
+from uuid import UUID
+
+import fsspec
+from fsspec.asyn import AsyncFileSystem
+from fsspec.utils import get_protocol
+from obstore.exceptions import GenericError
+from obstore.fsspec import register
+
+from union._initialize import get_storage
+from union._logging import logger
+
+
+def is_remote(path: typing.Union[pathlib.Path | str]) -> bool:
+    """
+    Let's find a replacement
+    """
+    protocol = get_protocol(str(path))
+    if protocol is None:
+        return False
+    return protocol != "file"
+
+
+def strip_file_header(path: str) -> str:
+    """
+    Drops file:// if it exists from the file
+    """
+    if path.startswith("file://"):
+        return path.replace("file://", "", 1)
+    return path
+
+
+def get_random_local_path(file_path_or_file_name: pathlib.Path | str | None = None) -> pathlib.Path:
+    """
+    Use file_path_or_file_name, when you want a random directory, but want to preserve the leaf file name
+    """
+    local_tmp = pathlib.Path(tempfile.mkdtemp(prefix="union-tmp-"))
+    key = UUID(int=random.getrandbits(128)).hex
+    tmp_folder = local_tmp / key
+    tail = ""
+    if file_path_or_file_name:
+        _, tail = os.path.split(file_path_or_file_name)
+    if tail:
+        tmp_folder.mkdir(parents=True, exist_ok=True)
+        return tmp_folder / tail
+    local_tmp.mkdir(parents=True, exist_ok=True)
+    return tmp_folder
+
+
+def get_random_local_directory() -> pathlib.Path:
+    """
+    :return: a random directory
+    :rtype: pathlib.Path
+    """
+    _dir = get_random_local_path(None)
+    pathlib.Path(_dir).mkdir(parents=True, exist_ok=True)
+    return _dir
+
+
+def get_underlying_filesystem(
+    protocol: typing.Optional[str] = None,
+    anonymous: bool = False,
+    path: typing.Optional[str] = None,
+    **kwargs,
+) -> fsspec.AbstractFileSystem:
+    if protocol is None:
+        # If protocol is None, get it from the path
+        protocol = get_protocol(path)
+
+    if get_storage():
+        storage_config = get_storage()
+        kwargs.update(storage_config.get_fsspec_kwargs(anonymous, **kwargs))
+    return fsspec.filesystem(protocol, **kwargs)
+
+
+def _get_anonymous_filesystem(from_path):
+    """Get the anonymous file system if needed."""
+    return get_underlying_filesystem(get_protocol(from_path), anonymous=True, asynchronous=True)
+
+
+async def get(from_path: str, to_path: Optional[str] = None, recursive: bool = False, **kwargs) -> str:
+    if not to_path:
+        name = pathlib.Path(from_path).name
+        to_path = get_random_local_path(file_path_or_file_name=name)
+        logger.debug(f"Storing file from {from_path} to {to_path}")
+    file_system = get_underlying_filesystem(path=from_path)
+    try:
+        return await _get_from_filesystem(file_system, from_path, to_path, recursive=recursive, **kwargs)
+    except (OSError, GenericError) as oe:
+        logger.debug(f"Error in getting {from_path} to {to_path} rec {recursive} {oe}")
+        if isinstance(file_system, AsyncFileSystem):
+            try:
+                exists = await file_system._exists(from_path)  # pylint: disable=W0212
+            except GenericError:
+                # for obstore, as it does not raise FileNotFoundError in fsspec but GenericError
+                # force it to try get_filesystem(anonymous=True)
+                exists = True
+        else:
+            exists = file_system.exists(from_path)
+        if not exists:
+            # TODO: update exception to be more specific
+            raise AssertionError(f"Unable to load data from {from_path}")
+        file_system = _get_anonymous_filesystem(from_path)
+        logger.debug(f"Attempting anonymous get with {file_system}")
+        return await _get_from_filesystem(file_system, from_path, to_path, recursive=recursive, **kwargs)
+
+
+async def _get_from_filesystem(
+    file_system: fsspec.AbstractFileSystem, from_path: str, to_path: pathlib.Path, recursive: bool, **kwargs
+):
+    if isinstance(file_system, AsyncFileSystem):
+        dst = await file_system._get(from_path, to_path, recursive=recursive, **kwargs)  # pylint: disable=W0212
+    else:
+        dst = file_system.get(from_path, to_path, recursive=recursive, **kwargs)
+
+    if isinstance(dst, (str, pathlib.Path)):
+        return dst
+    return to_path
+
+
+async def put(from_path: str, to_path: Optional[str] = None, recursive: bool = False, **kwargs):
+    if not to_path:
+        from union._context import internal_ctx
+
+        ctx = internal_ctx()
+        name = pathlib.Path(from_path).name if not recursive else None  # don't pass a name for folders
+        to_path = ctx.raw_data.get_random_remote_path(file_name=name)
+
+    file_system = get_underlying_filesystem(path=to_path)
+    from_path = strip_file_header(from_path)
+    if isinstance(file_system, AsyncFileSystem):
+        dst = await file_system._put(from_path, to_path, recursive=recursive, **kwargs)  # pylint: disable=W0212
+    else:
+        dst = file_system.put(from_path, to_path, recursive=recursive, **kwargs)
+    if isinstance(dst, (str, pathlib.Path)):
+        return dst
+    else:
+        return to_path
+
+
+async def put_stream(
+    data_iterable: typing.AsyncIterable[bytes] | bytes, *, name: str | None = None, to_path: str | None = None, **kwargs
+) -> str:
+    """
+    Put a stream of data to a remote location. This is useful for streaming data to a remote location.
+    Example usage:
+    ```python
+    import union.storage as storage
+    storage.put_stream(iter([b'hello']), name="my_file.txt")
+    OR
+    storage.put_stream(iter([b'hello']), to_path="s3://my_bucket/my_file.txt")
+    ```
+
+    :param data_iterable: Iterable of bytes to be streamed.
+    :param name: Name of the file to be created. If not provided, a random name will be generated.
+    :param to_path: Path to the remote location where the data will be stored.
+    :param kwargs: Additional arguments to be passed to the underlying filesystem.
+    :rtype: str
+    :return: The path to the remote location where the data was stored.
+    """
+    if not to_path:
+        from union._context import internal_ctx
+
+        ctx = internal_ctx()
+        to_path = ctx.raw_data.get_random_remote_path(file_name=name)
+    fs = get_underlying_filesystem(path=to_path)
+    file_handle = None
+    if isinstance(fs, AsyncFileSystem):
+        try:
+            file_handle = await fs.open_async(to_path, "wb", **kwargs)
+            if isinstance(data_iterable, bytes):
+                await file_handle.write(data_iterable)
+            else:
+                async for data in data_iterable:
+                    await file_handle.write(data)
+            return
+        except NotImplementedError:
+            logger.debug(f"{fs} doesn't implement 'open_async', falling back to sync")
+        finally:
+            if file_handle is not None:
+                await file_handle.close()
+
+    with fs.open(to_path, "wb", **kwargs) as f:
+        if isinstance(data_iterable, bytes):
+            f.write(data_iterable)
+        else:
+            # If data_iterable is async iterable, iterate over it and write each chunk to the file
+            async for data in data_iterable:
+                f.write(data)
+    return str(to_path)
+
+
+async def get_stream(path: pathlib.Path, chunk_size=10 * 2**20, **kwargs) -> AsyncIterator[bytes]:
+    """
+    Get a stream of data from a remote location.
+    This is useful for downloading streaming data from a remote location.
+    Example usage:
+    ```python
+    import union.storage as storage
+    obj = storage.get_stream(path="s3://my_bucket/my_file.txt")
+    ```
+
+    :param path: Path to the remote location where the data will be downloaded.
+    :param kwargs: Additional arguments to be passed to the underlying filesystem.
+    :param chunk_size: Size of each chunk to be read from the file.
+    :return: An async iterator that yields chunks of data.
+    """
+    fs = get_underlying_filesystem(path=path, **kwargs)
+    file_size = fs.info(path)["size"]
+    total_read = 0
+    file_handle = None
+    try:
+        if isinstance(fs, AsyncFileSystem):
+            file_handle = await fs.open_async(path, "rb")
+            while chunk := await file_handle.read(min(chunk_size, file_size - total_read)):
+                total_read += len(chunk)
+                yield chunk
+            return
+    except NotImplementedError:
+        logger.debug(f"{fs} doesn't implement 'open_async', falling back to sync")
+    finally:
+        if file_handle is not None:
+            file_handle.close()
+
+    # Sync fallback
+    with fs.open(path, "rb") as file_handle:
+        while chunk := file_handle.read(min(chunk_size, file_size - total_read)):
+            total_read += len(chunk)
+            yield chunk
+
+
+def join(*paths: str) -> str:
+    """
+    Join multiple paths together. This is a wrapper around os.path.join.
+    # TODO replace with proper join with fsspec root etc
+
+    :param paths: Paths to be joined.
+    """
+    return str(os.path.join(*paths))
+
+
+register(["s3", "gs", "abfs", "abfss"], asynchronous=True)

union/storage/_utils.py

@@ -0,0 +1,5 @@
+import os
+
+# This is the default chunk size union will use for writing to S3 and GCS. This is set to 25MB by default and is
+# configurable by the user if needed. This is used when put() is called on filesystems.
+_WRITE_SIZE_CHUNK_BYTES = int(os.environ.get("_F_P_WRITE_CHUNK_SIZE", "26214400"))  # 25 * 2**20

union/types/__init__.py

@@ -0,0 +1,11 @@
+from ._renderer import Renderable
+from ._string_literals import literal_string_repr
+from ._type_engine import TypeEngine, TypeTransformer, TypeTransformerFailedError
+
+__all__ = [
+    "Renderable",
+    "TypeEngine",
+    "TypeTransformer",
+    "TypeTransformerFailedError",
+    "literal_string_repr",
+]

union/types/_renderer.py

@@ -0,0 +1,162 @@
+from typing import TYPE_CHECKING, Any
+
+from markdown_it import MarkdownIt
+from typing_extensions import Protocol, runtime_checkable
+
+from union._utils import lazy_module
+
+if TYPE_CHECKING:
+    # Always import these modules in type-checking mode or when running pytest
+    import pandas
+    import pyarrow
+else:
+    pandas = lazy_module("pandas")
+    pyarrow = lazy_module("pyarrow")
+
+
+@runtime_checkable
+class Renderable(Protocol):
+    def to_html(self, python_value: Any) -> str:
+        """Convert an object(markdown, pandas.dataframe) to HTML and return HTML as a unicode string.
+        Returns: An HTML document as a string.
+        """
+        raise NotImplementedError
+
+
+DEFAULT_MAX_ROWS = 10
+DEFAULT_MAX_COLS = 100
+
+
+class TopFrameRenderer:
+    """
+    Render a DataFrame as an HTML table.
+    """
+
+    def __init__(self, max_rows: int = DEFAULT_MAX_ROWS, max_cols: int = DEFAULT_MAX_COLS):
+        self._max_rows = max_rows
+        self._max_cols = max_cols
+
+    def to_html(self, df: "pandas.DataFrame") -> str:
+        assert isinstance(df, pandas.DataFrame)
+        return df.to_html(max_rows=self._max_rows, max_cols=self._max_cols)
+
+
+class ArrowRenderer:
+    """
+    Render an Arrow dataframe as an HTML table.
+    """
+
+    def to_html(self, df: "pyarrow.Table") -> str:
+        assert isinstance(df, pyarrow.Table)
+        return df.to_string()
+
+
+class MarkdownRenderer:
+    """Convert a markdown string to HTML and return HTML as a unicode string."""
+
+    def to_html(self, text: str) -> str:
+        return MarkdownIt().render(text)
+
+
+class SourceCodeRenderer:
+    """
+    Convert Python source code to HTML, and return HTML as a unicode string.
+    """
+
+    def __init__(self, title: str = "Source Code"):
+        self._title = title
+
+    def to_html(self, source_code: str) -> str:
+        """
+        Convert the provided Python source code into HTML format using Pygments library.
+
+        This method applies a colorful style and replaces the color "#fff0f0" with "#ffffff" in CSS.
+
+        Args:
+            source_code (str): The Python source code to be converted.
+
+        Returns:
+            str: The resulting HTML as a string, including CSS and highlighted source code.
+        """
+        from pygments import highlight
+        from pygments.formatters.html import HtmlFormatter
+        from pygments.lexers.python import PythonLexer
+
+        formatter = HtmlFormatter(style="colorful")
+        css = formatter.get_style_defs(".highlight").replace("#fff0f0", "#ffffff")
+        html = highlight(source_code, PythonLexer(), formatter)
+        return f"<style>{css}</style>{html}"
+
+
+class PythonDependencyRenderer:
+    """
+    PythonDependencyDeck is a deck that contains information about packages installed via pip.
+    """
+
+    def __init__(self, title: str = "Dependencies"):
+        self._title = title
+
+    def to_html(self) -> str:
+        import json
+        import subprocess
+        import sys
+
+        from flytekit.loggers import logger
+
+        try:
+            installed_packages = json.loads(
+                subprocess.check_output([sys.executable, "-m", "pip", "list", "--format", "json"])
+            )
+            requirements_txt = (
+                subprocess.check_output([sys.executable, "-m", "pip", "freeze"])
+                .decode("utf-8")
+                .replace("\\n", "\n")
+                .rstrip()
+            )
+        except Exception as e:
+            logger.error(f"Error occurred while fetching installed packages: {e}")
+            return "Error occurred while fetching installed packages."
+
+        table = (
+            "<table>\n<tr>\n<th style='text-align:left;'>Name</th>\n<th style='text-align:left;'>Version</th>\n</tr>\n"
+        )
+
+        for entry in installed_packages:
+            table += f"<tr>\n<td>{entry['name']}</td>\n<td>{entry['version']}</td>\n</tr>\n"
+
+        table += "</table>"
+
+        html = f"""
+        <!DOCTYPE html>
+        <html lang="en">
+        <head>
+        <meta charset="UTF-8">
+        <meta name="viewport" content="width=device-width, initial-scale=1.0">
+        <title>Flyte Dependencies</title>
+        <script>
+        async function copyTable() {{
+          var requirements_txt = document.getElementById('requirements_txt');
+
+          try {{
+            await navigator.clipboard.writeText(requirements_txt.innerText);
+          }} catch (err) {{
+            console.log('Error accessing the clipboard: ' + err);
+          }}
+        }}
+        </script>
+        </head>
+        <body>
+
+        <button onclick="copyTable()">
+          <span>Copy table as requirements.txt</span>
+        </button>
+        <h3>Python Dependencies</h3>
+
+        {table}
+
+        <div id="requirements_txt" style="display:none">{requirements_txt}</div>
+
+        </body>
+        </html>
+        """
+        return html

union/types/_string_literals.py

@@ -0,0 +1,120 @@
+import base64
+import json
+from typing import Any, Dict, Union
+
+import msgpack
+from flyteidl.core import literals_pb2
+from google.protobuf.json_format import MessageToDict
+
+from union._protos.workflow import run_definition_pb2
+
+
+def _primitive_to_string(primitive: literals_pb2.Primitive) -> Any:
+    """
+    This method is used to convert a primitive to a string representation.
+    """
+    match primitive.WhichOneof("value"):
+        case "integer":
+            return primitive.integer
+        case "float_value":
+            return primitive.float_value
+        case "boolean":
+            return primitive.boolean
+        case "string_value":
+            return primitive.string_value
+        case "datetime":
+            return primitive.datetime.ToDatetime().isoformat()
+        case "duration":
+            return primitive.duration.ToSeconds()
+        case _:
+            raise ValueError(f"Unknown primitive type {primitive}")
+
+
+def _scalar_to_string(scalar: literals_pb2.Scalar) -> Any:
+    """
+    This method is used to convert a scalar to a string representation.
+    """
+    match scalar.WhichOneof("value"):
+        case "primitive":
+            return _primitive_to_string(scalar.primitive)
+        case "none_type":
+            return None
+        case "error":
+            return scalar.error.message
+        case "structured_dataset":
+            return scalar.structured_dataset.uri
+        case "schema":
+            return scalar.schema.uri
+        case "blob":
+            return scalar.blob.uri
+        case "binary":
+            if scalar.binary.tag == "msgpack":
+                return json.dumps(msgpack.unpackb(scalar.binary.value))
+            return base64.b64encode(scalar.binary.value)
+        case "generic":
+            return MessageToDict(scalar.generic)
+        case "union":
+            return _literal_string_repr(scalar.union.value)
+        case _:
+            raise ValueError(f"Unknown scalar type {scalar}")
+
+
+def _literal_string_repr(lit: literals_pb2.Literal) -> Any:
+    """
+    This method is used to convert a literal to a string representation. This is useful in places, where we need to
+    use a shortened string representation of a literal, especially a FlyteFile, FlyteDirectory, or StructuredDataset.
+    """
+    match lit.WhichOneof("value"):
+        case "scalar":
+            return _scalar_to_string(lit.scalar)
+        case "collection":
+            return [literal_string_repr(i) for i in lit.collection.literals]
+        case "map":
+            return {k: literal_string_repr(v) for k, v in lit.map.literals.items()}
+        case "offloaded_metadata":
+            # TODO: load literal from offloaded literal?
+            return f"Offloaded literal metadata: {lit.offloaded_metadata}"
+        case _:
+            raise ValueError(f"Unknown literal type {lit}")
+
+
+def _dict_literal_repr(lmd: Dict[str, literals_pb2.Literal]) -> Dict[str, Any]:
+    """
+    This method is used to convert a literal map to a string representation.
+    """
+    return {k: _literal_string_repr(v) for k, v in lmd.items()}
+
+
+def literal_string_repr(
+    lm: Union[
+        literals_pb2.Literal,
+        run_definition_pb2.NamedLiteral,
+        run_definition_pb2.Inputs,
+        run_definition_pb2.Outputs,
+        literals_pb2.LiteralMap,
+        Dict[str, literals_pb2.Literal],
+    ],
+) -> Dict[str, Any]:
+    """
+    This method is used to convert a literal map to a string representation.
+    """
+    if lm is None:
+        return {}
+    match lm:
+        case literals_pb2.Literal():
+            return _literal_string_repr(lm)
+        case literals_pb2.LiteralMap():
+            return _dict_literal_repr(lm.literals)
+        case run_definition_pb2.NamedLiteral():
+            lmd = {lm.name: lm.value}
+            return _dict_literal_repr(lmd)
+        case run_definition_pb2.Inputs():
+            lmd = {n.name: n.value for n in lm.literals}
+            return _dict_literal_repr(lmd)
+        case run_definition_pb2.Outputs():
+            lmd = {n.name: n.value for n in lm.literals}
+            return _dict_literal_repr(lmd)
+        case dict():
+            return _dict_literal_repr(lm)
+        case _:
+            raise ValueError(f"Unknown literal type {lm}, type{type(lm)}")