flyte 0.0.1b0__py3-none-any.whl → 2.0.0b46__py3-none-any.whl

flyte/storage/_config.py

@@ -0,0 +1,237 @@
+from __future__ import annotations
+
+import datetime
+import os
+import typing
+from dataclasses import dataclass
+from typing import ClassVar
+
+from flyte.config import set_if_exists
+
+
+@dataclass(init=True, repr=True, eq=True, frozen=True)
+class Storage(object):
+    """
+    Data storage configuration that applies across any provider.
+    """
+
+    retries: int = 3
+    backoff: datetime.timedelta = datetime.timedelta(seconds=5)
+    enable_debug: bool = False
+    attach_execution_metadata: bool = True
+
+    _KEY_ENV_VAR_MAPPING: ClassVar[typing.Dict[str, str]] = {
+        "enable_debug": "UNION_STORAGE_DEBUG",
+        "retries": "UNION_STORAGE_RETRIES",
+        "backoff": "UNION_STORAGE_BACKOFF_SECONDS",
+    }
+
+    def get_fsspec_kwargs(self, anonymous: bool = False, **kwargs) -> typing.Dict[str, typing.Any]:
+        """
+        Returns the configuration as kwargs for constructing an fsspec filesystem.
+        """
+        return {}
+
+    @classmethod
+    def _auto_as_kwargs(cls) -> typing.Dict[str, typing.Any]:
+        retries = os.getenv(cls._KEY_ENV_VAR_MAPPING["retries"])
+        backoff = os.getenv(cls._KEY_ENV_VAR_MAPPING["backoff"])
+        enable_debug = os.getenv(cls._KEY_ENV_VAR_MAPPING["enable_debug"])
+
+        kwargs: typing.Dict[str, typing.Any] = {}
+        kwargs = set_if_exists(kwargs, "enable_debug", enable_debug)
+        kwargs = set_if_exists(kwargs, "retries", retries)
+        kwargs = set_if_exists(kwargs, "backoff", backoff)
+        return kwargs
+
+    @classmethod
+    def auto(cls) -> Storage:
+        """
+        Construct the config object automatically from environment variables.
+        """
+        return cls(**cls._auto_as_kwargs())
+
+
+@dataclass(init=True, repr=True, eq=True, frozen=True)
+class S3(Storage):
+    """
+    S3 specific configuration
+    """
+
+    endpoint: typing.Optional[str] = None
+    access_key_id: typing.Optional[str] = None
+    secret_access_key: typing.Optional[str] = None
+    region: typing.Optional[str] = None
+
+    _KEY_ENV_VAR_MAPPING: ClassVar[typing.Dict[str, str]] = {
+        "endpoint": "FLYTE_AWS_ENDPOINT",
+        "access_key_id": "FLYTE_AWS_ACCESS_KEY_ID",
+        "secret_access_key": "FLYTE_AWS_SECRET_ACCESS_KEY",
+    } | Storage._KEY_ENV_VAR_MAPPING
+
+    # Refer to https://github.com/developmentseed/obstore/blob/33654fc37f19a657689eb93327b621e9f9e01494/obstore/python/obstore/store/_aws.pyi#L11
+    # for key and secret
+    _CONFIG_KEY_FSSPEC_S3_KEY_ID: ClassVar = "access_key_id"
+    _CONFIG_KEY_FSSPEC_S3_SECRET: ClassVar = "secret_access_key"
+    _CONFIG_KEY_ENDPOINT: ClassVar = "endpoint_url"
+    _KEY_SKIP_SIGNATURE: ClassVar = "skip_signature"
+
+    @classmethod
+    def auto(cls, region: str | None = None) -> S3:
+        """
+        :return: Config
+        """
+        endpoint = os.getenv(cls._KEY_ENV_VAR_MAPPING["endpoint"], None)
+        access_key_id = os.getenv(cls._KEY_ENV_VAR_MAPPING["access_key_id"], None)
+        secret_access_key = os.getenv(cls._KEY_ENV_VAR_MAPPING["secret_access_key"], None)
+
+        kwargs = super()._auto_as_kwargs()
+        kwargs = set_if_exists(kwargs, "endpoint", endpoint)
+        kwargs = set_if_exists(kwargs, "access_key_id", access_key_id)
+        kwargs = set_if_exists(kwargs, "secret_access_key", secret_access_key)
+        kwargs = set_if_exists(kwargs, "region", region)
+
+        return S3(**kwargs)
+
+    @classmethod
+    def for_sandbox(cls) -> S3:
+        """
+        :return:
+        """
+        kwargs = super()._auto_as_kwargs()
+        final_kwargs = kwargs | {
+            "endpoint": "http://localhost:4566",
+            "access_key_id": "minio",
+            "secret_access_key": "miniostorage",
+        }
+        return S3(**final_kwargs)
+
+    def get_fsspec_kwargs(self, anonymous: bool = False, **kwargs) -> typing.Dict[str, typing.Any]:
+        # Construct the config object
+        kwargs.pop("anonymous", None)  # Remove anonymous if it exists, as we handle it separately
+        config: typing.Dict[str, typing.Any] = {}
+        if self._CONFIG_KEY_FSSPEC_S3_KEY_ID in kwargs or self.access_key_id:
+            config[self._CONFIG_KEY_FSSPEC_S3_KEY_ID] = kwargs.pop(
+                self._CONFIG_KEY_FSSPEC_S3_KEY_ID, self.access_key_id
+            )
+        if self._CONFIG_KEY_FSSPEC_S3_SECRET in kwargs or self.secret_access_key:
+            config[self._CONFIG_KEY_FSSPEC_S3_SECRET] = kwargs.pop(
+                self._CONFIG_KEY_FSSPEC_S3_SECRET, self.secret_access_key
+            )
+        if self._CONFIG_KEY_ENDPOINT in kwargs or self.endpoint:
+            config["endpoint_url"] = kwargs.pop(self._CONFIG_KEY_ENDPOINT, self.endpoint)
+
+        retries = kwargs.pop("retries", self.retries)
+        backoff = kwargs.pop("backoff", self.backoff)
+
+        if anonymous:
+            config[self._KEY_SKIP_SIGNATURE] = True
+
+        retry_config = {
+            "max_retries": retries,
+            "backoff": {
+                "base": 2,
+                "init_backoff": backoff,
+                "max_backoff": datetime.timedelta(seconds=16),
+            },
+            "retry_timeout": datetime.timedelta(minutes=3),
+        }
+
+        client_options = {"timeout": "99999s", "allow_http": True}
+
+        if config:
+            kwargs["config"] = config
+        kwargs["client_options"] = client_options or None
+        kwargs["retry_config"] = retry_config or None
+        if self.region:
+            kwargs["region"] = self.region
+
+        return kwargs
+
+
+@dataclass(init=True, repr=True, eq=True, frozen=True)
+class GCS(Storage):
+    """
+    Any GCS specific configuration.
+    """
+
+    gsutil_parallelism: bool = False
+
+    _KEY_ENV_VAR_MAPPING: ClassVar[dict[str, str]] = {
+        "gsutil_parallelism": "GCP_GSUTIL_PARALLELISM",
+    }
+
+    @classmethod
+    def auto(cls) -> GCS:
+        gsutil_parallelism = os.getenv(cls._KEY_ENV_VAR_MAPPING["gsutil_parallelism"], None)
+
+        kwargs: typing.Dict[str, typing.Any] = {}
+        kwargs = set_if_exists(kwargs, "gsutil_parallelism", gsutil_parallelism)
+        return GCS(**kwargs)
+
+    def get_fsspec_kwargs(self, anonymous: bool = False, **kwargs) -> typing.Dict[str, typing.Any]:
+        kwargs.pop("anonymous", None)
+        return kwargs
+
+
+@dataclass(init=True, repr=True, eq=True, frozen=True)
+class ABFS(Storage):
+    """
+    Any Azure Blob Storage specific configuration.
+    """
+
+    account_name: typing.Optional[str] = None
+    account_key: typing.Optional[str] = None
+    tenant_id: typing.Optional[str] = None
+    client_id: typing.Optional[str] = None
+    client_secret: typing.Optional[str] = None
+
+    _KEY_ENV_VAR_MAPPING: ClassVar[dict[str, str]] = {
+        "account_name": "AZURE_STORAGE_ACCOUNT_NAME",
+        "account_key": "AZURE_STORAGE_ACCOUNT_KEY",
+        "tenant_id": "AZURE_TENANT_ID",
+        "client_id": "AZURE_CLIENT_ID",
+        "client_secret": "AZURE_CLIENT_SECRET",
+    }
+    _KEY_SKIP_SIGNATURE: ClassVar = "skip_signature"
+
+    @classmethod
+    def auto(cls) -> ABFS:
+        account_name = os.getenv(cls._KEY_ENV_VAR_MAPPING["account_name"], None)
+        account_key = os.getenv(cls._KEY_ENV_VAR_MAPPING["account_key"], None)
+        tenant_id = os.getenv(cls._KEY_ENV_VAR_MAPPING["tenant_id"], None)
+        client_id = os.getenv(cls._KEY_ENV_VAR_MAPPING["client_id"], None)
+        client_secret = os.getenv(cls._KEY_ENV_VAR_MAPPING["client_secret"], None)
+
+        kwargs: typing.Dict[str, typing.Any] = {}
+        kwargs = set_if_exists(kwargs, "account_name", account_name)
+        kwargs = set_if_exists(kwargs, "account_key", account_key)
+        kwargs = set_if_exists(kwargs, "tenant_id", tenant_id)
+        kwargs = set_if_exists(kwargs, "client_id", client_id)
+        kwargs = set_if_exists(kwargs, "client_secret", client_secret)
+        return ABFS(**kwargs)
+
+    def get_fsspec_kwargs(self, anonymous: bool = False, **kwargs) -> typing.Dict[str, typing.Any]:
+        kwargs.pop("anonymous", None)
+        config: typing.Dict[str, typing.Any] = {}
+        if "account_name" in kwargs or self.account_name:
+            config["account_name"] = kwargs.get("account_name", self.account_name)
+        if "account_key" in kwargs or self.account_key:
+            config["account_key"] = kwargs.get("account_key", self.account_key)
+        if "client_id" in kwargs or self.client_id:
+            config["client_id"] = kwargs.get("client_id", self.client_id)
+        if "client_secret" in kwargs or self.client_secret:
+            config["client_secret"] = kwargs.get("client_secret", self.client_secret)
+        if "tenant_id" in kwargs or self.tenant_id:
+            config["tenant_id"] = kwargs.get("tenant_id", self.tenant_id)
+
+        if anonymous:
+            config[self._KEY_SKIP_SIGNATURE] = True
+
+        client_options = {"timeout": "99999s", "allow_http": "true"}
+
+        if config:
+            kwargs["config"] = config
+        kwargs["client_options"] = client_options
+
+        return kwargs

flyte/storage/_parallel_reader.py

@@ -0,0 +1,289 @@
+from __future__ import annotations
+
+import asyncio
+import dataclasses
+import io
+import os
+import pathlib
+import sys
+import tempfile
+import typing
+from typing import Any, Hashable, Protocol
+
+import aiofiles
+import aiofiles.os
+import obstore
+
+if typing.TYPE_CHECKING:
+    from obstore import Bytes, ObjectMeta
+    from obstore.store import ObjectStore
+
+CHUNK_SIZE = int(os.getenv("FLYTE_IO_CHUNK_SIZE", str(16 * 1024 * 1024)))
+MAX_CONCURRENCY = int(os.getenv("FLYTE_IO_MAX_CONCURRENCY", str(32)))
+
+
+class DownloadQueueEmpty(RuntimeError):
+    pass
+
+
+class BufferProtocol(Protocol):
+    async def write(self, offset, length, value: Bytes) -> None: ...
+
+    async def read(self) -> memoryview: ...
+
+    @property
+    def complete(self) -> bool: ...
+
+
+@dataclasses.dataclass
+class _MemoryBuffer:
+    arr: bytearray
+    pending: int
+    _closed: bool = False
+
+    async def write(self, offset: int, length: int, value: Bytes) -> None:
+        self.arr[offset : offset + length] = memoryview(value)
+        self.pending -= length
+
+    async def read(self) -> memoryview:
+        return memoryview(self.arr)
+
+    @property
+    def complete(self) -> bool:
+        return self.pending == 0
+
+    @classmethod
+    def new(cls, size):
+        return cls(arr=bytearray(size), pending=size)
+
+
+@dataclasses.dataclass
+class _FileBuffer:
+    path: pathlib.Path
+    pending: int
+    _handle: io.FileIO | None = None
+    _closed: bool = False
+
+    async def write(self, offset: int, length: int, value: Bytes) -> None:
+        async with aiofiles.open(self.path, mode="r+b") as f:
+            await f.seek(offset)
+            await f.write(value)
+        self.pending -= length
+
+    async def read(self) -> memoryview:
+        async with aiofiles.open(self.path, mode="rb") as f:
+            return memoryview(await f.read())
+
+    @property
+    def complete(self) -> bool:
+        return self.pending == 0
+
+    @classmethod
+    def new(cls, path: pathlib.Path, size: int):
+        path.parent.mkdir(parents=True, exist_ok=True)
+        path.touch()
+        return cls(path=path, pending=size)
+
+
+@dataclasses.dataclass
+class Chunk:
+    offset: int
+    length: int
+
+
+@dataclasses.dataclass
+class Source:
+    id: Hashable
+    path: pathlib.Path  # Should be str, represents the fully qualified prefix of a file (no bucket)
+    length: int
+    offset: int = 0
+    metadata: Any | None = None
+
+
+@dataclasses.dataclass
+class DownloadTask:
+    source: Source
+    chunk: Chunk
+    target: pathlib.Path | None = None
+
+
+class ObstoreParallelReader:
+    def __init__(
+        self,
+        store: ObjectStore,
+        *,
+        chunk_size=CHUNK_SIZE,
+        max_concurrency=MAX_CONCURRENCY,
+    ):
+        self._store = store
+        self._chunk_size = chunk_size
+        self._max_concurrency = max_concurrency
+
+    def _chunks(self, size) -> typing.Iterator[tuple[int, int]]:
+        cs = self._chunk_size
+        for offset in range(0, size, cs):
+            length = min(cs, size - offset)
+            yield offset, length
+
+    async def _as_completed(self, gen: typing.AsyncGenerator[DownloadTask, None], transformer=None):
+        inq: asyncio.Queue = asyncio.Queue(self._max_concurrency * 2)
+        outq: asyncio.Queue = asyncio.Queue()
+        sentinel = object()
+        done = asyncio.Event()
+
+        active: dict[Hashable, _FileBuffer | _MemoryBuffer] = {}
+
+        async def _fill():
+            # Helper function to fill the input queue, this is because the generator is async because it does list/head
+            # calls on the object store which are async.
+            try:
+                counter = 0
+                async for task in gen:
+                    if task.source.id not in active:
+                        active[task.source.id] = (
+                            _FileBuffer.new(task.target, task.source.length)
+                            if task.target is not None
+                            else _MemoryBuffer.new(task.source.length)
+                        )
+                    await inq.put(task)
+                    counter += 1
+                await inq.put(sentinel)
+                if counter == 0:
+                    raise DownloadQueueEmpty
+            except asyncio.CancelledError:
+                # document why we need to swallow this
+                pass
+
+        async def _worker():
+            try:
+                while not done.is_set():
+                    task: DownloadTask = await inq.get()
+                    if task is sentinel:
+                        inq.put_nowait(sentinel)
+                        break
+                    # chunk.offset is the local offset within the source (e.g., 0, chunk_size, 2*chunk_size)
+                    # source.offset is the offset within the file where the source data starts
+                    # The actual file position is the sum of both
+                    file_offset = task.chunk.offset + task.source.offset
+                    buf = active[task.source.id]
+                    data_to_write = await obstore.get_range_async(
+                        self._store,
+                        str(task.source.path),
+                        start=file_offset,
+                        end=file_offset + task.chunk.length,
+                    )
+                    await buf.write(
+                        task.chunk.offset,
+                        task.chunk.length,
+                        data_to_write,
+                    )
+                    if not buf.complete:
+                        continue
+                    if transformer is not None:
+                        result = await transformer(buf, task.source)
+                    elif task.target is not None:
+                        result = task.target
+                    else:
+                        result = task.source
+                    outq.put_nowait((task.source.id, result))
+                    del active[task.source.id]
+            except asyncio.CancelledError:
+                pass
+            finally:
+                done.set()
+
+        # Yield results as they are completed
+        if sys.version_info >= (3, 11):
+            async with asyncio.TaskGroup() as tg:
+                tg.create_task(_fill())
+                for _ in range(self._max_concurrency):
+                    tg.create_task(_worker())
+                while not done.is_set():
+                    yield await outq.get()
+        else:
+            fill_task = asyncio.create_task(_fill())
+            worker_tasks = [asyncio.create_task(_worker()) for _ in range(self._max_concurrency)]
+            try:
+                while not done.is_set():
+                    yield await outq.get()
+            except Exception as e:
+                if not fill_task.done():
+                    fill_task.cancel()
+                for wt in worker_tasks:
+                    if not wt.done():
+                        wt.cancel()
+                raise e
+            finally:
+                await asyncio.gather(fill_task, *worker_tasks, return_exceptions=True)
+
+        # Drain the output queue
+        try:
+            while True:
+                yield outq.get_nowait()
+        except asyncio.QueueEmpty:
+            pass
+
+    async def download_files(
+        self,
+        src_prefix: pathlib.Path,
+        target_prefix: pathlib.Path,
+        *paths,
+        destination_file_name: str | None = None,
+        exclude: list[str] | None = None,
+    ) -> None:
+        """
+        src_prefix: Prefix you want to download from in the object store, not including the bucket name, nor file name.
+                    Should be replaced with string
+        target_prefix: Local directory to download to
+        paths: Specific paths (relative to src_prefix) to download. If empty, download everything
+        exclude: List of patterns to exclude from the download.
+        """
+
+        def _keep(path):
+            if exclude is not None and any(path.match(e) for e in exclude):
+                return False
+            return True
+
+        async def _list_downloadable() -> typing.AsyncGenerator[ObjectMeta, None]:
+            if paths:
+                # For specific file paths, use async head
+                for path_ in paths:
+                    path = src_prefix / path_
+                    if _keep(path):
+                        yield await obstore.head_async(self._store, str(path))
+                return
+
+            # Use obstore.list() for recursive listing (all files in all subdirectories)
+            # obstore.list() returns an async iterator that yields batches (lists) of objects
+            async for batch in obstore.list(self._store, prefix=str(src_prefix)):
+                for obj in batch:
+                    yield obj
+
+        async def _gen(tmp_dir: str) -> typing.AsyncGenerator[DownloadTask, None]:
+            async for obj in _list_downloadable():
+                path = pathlib.Path(obj["path"])  # e.g. Path(prefix/file.txt), needs to be changed to str.
+                size = obj["size"]
+                source = Source(id=path, path=path, length=size)
+                # Strip src_prefix from path for destination
+                rel_path = path.relative_to(src_prefix)  # doesn't work on windows
+                for offset, length in self._chunks(size):
+                    yield DownloadTask(
+                        source=source,
+                        target=tmp_dir / rel_path,  # doesn't work on windows
+                        chunk=Chunk(offset, length),
+                    )
+
+        def _transform_decorator(tmp_dir: str):
+            async def _transformer(buf: _FileBuffer, _: Source) -> None:
+                if len(paths) == 1 and destination_file_name is not None:
+                    target = target_prefix / destination_file_name
+                else:
+                    target = target_prefix / buf.path.relative_to(tmp_dir)
+                await aiofiles.os.makedirs(target.parent, exist_ok=True)
+                return await aiofiles.os.replace(buf.path, target)  # mv buf.path target
+
+            return _transformer
+
+        with tempfile.TemporaryDirectory() as temporary_dir:
+            async for _ in self._as_completed(_gen(temporary_dir), transformer=_transform_decorator(temporary_dir)):
+                pass