flyte 0.2.0b1__py3-none-any.whl → 2.0.0b46__py3-none-any.whl

flyte/storage/_storage.py

@@ -1,19 +1,38 @@
+from __future__ import annotations
+
 import os
 import pathlib
 import random
 import tempfile
 import typing
-from typing import AsyncIterator, Optional
+from typing import AsyncGenerator, Optional
 from uuid import UUID
 
 import fsspec
+import obstore
 from fsspec.asyn import AsyncFileSystem
 from fsspec.utils import get_protocol
 from obstore.exceptions import GenericError
 from obstore.fsspec import register
+from obstore.store import ObjectStore
 
 from flyte._initialize import get_storage
 from flyte._logging import logger
+from flyte.errors import InitializationError, OnlyAsyncIOSupportedError
+
+if typing.TYPE_CHECKING:
+    from obstore import AsyncReadableFile, AsyncWritableFile
+
+_OBSTORE_SUPPORTED_PROTOCOLS = ["s3", "gs", "abfs", "abfss"]
+
+
+def _is_obstore_supported_protocol(protocol: str) -> bool:
+    """
+    Check if the given protocol is supported by obstore.
+    :param protocol: Protocol to check.
+    :return: True if the protocol is supported, False otherwise.
+    """
+    return protocol in _OBSTORE_SUPPORTED_PROTOCOLS
 
 
 def is_remote(path: typing.Union[pathlib.Path | str]) -> bool:
@@ -62,6 +81,52 @@ def get_random_local_directory() -> pathlib.Path:
     return _dir
 
 
+def get_configured_fsspec_kwargs(
+    protocol: typing.Optional[str] = None, anonymous: bool = False
+) -> typing.Dict[str, typing.Any]:
+    if protocol:
+        # Try to get storage config safely - may not be initialized for local operations
+        try:
+            storage_config = get_storage()
+        except InitializationError:
+            storage_config = None
+
+        match protocol:
+            case "s3":
+                # If the protocol is s3, we can use the s3 filesystem
+                from flyte.storage import S3
+
+                if storage_config and isinstance(storage_config, S3):
+                    return storage_config.get_fsspec_kwargs(anonymous=anonymous)
+
+                return S3.auto().get_fsspec_kwargs(anonymous=anonymous)
+            case "gs":
+                # If the protocol is gs, we can use the gs filesystem
+                from flyte.storage import GCS
+
+                if storage_config and isinstance(storage_config, GCS):
+                    return storage_config.get_fsspec_kwargs(anonymous=anonymous)
+
+                return GCS.auto().get_fsspec_kwargs(anonymous=anonymous)
+            case "abfs" | "abfss":
+                # If the protocol is abfs or abfss, we can use the abfs filesystem
+                from flyte.storage import ABFS
+
+                if storage_config and isinstance(storage_config, ABFS):
+                    return storage_config.get_fsspec_kwargs(anonymous=anonymous)
+
+                return ABFS.auto().get_fsspec_kwargs(anonymous=anonymous)
+            case _:
+                return {}
+
+    # If no protocol, return args from storage config if set
+    storage_config = get_storage()
+    if storage_config:
+        return storage_config.get_fsspec_kwargs(anonymous)
+
+    return {}
+
+
 def get_underlying_filesystem(
     protocol: typing.Optional[str] = None,
     anonymous: bool = False,
@@ -72,10 +137,10 @@ def get_underlying_filesystem(
         # If protocol is None, get it from the path
         protocol = get_protocol(path)
 
-    storage_config = get_storage()
-    if storage_config:
-        kwargs.update(storage_config.get_fsspec_kwargs(anonymous, **kwargs))
-    return fsspec.filesystem(protocol, **kwargs)
+    configured_kwargs = get_configured_fsspec_kwargs(protocol, anonymous=anonymous)
+    configured_kwargs.update(kwargs)
+
+    return fsspec.filesystem(protocol, **configured_kwargs)
 
 
 def _get_anonymous_filesystem(from_path):
@@ -83,16 +148,86 @@ def _get_anonymous_filesystem(from_path):
     return get_underlying_filesystem(get_protocol(from_path), anonymous=True, asynchronous=True)
 
 
+async def _get_obstore_bypass(
+    from_path: str,
+    to_path: str | pathlib.Path,
+    recursive: bool = False,
+    exclude: list[str] | None = None,
+    **kwargs,
+) -> str:
+    from flyte.storage._parallel_reader import ObstoreParallelReader
+
+    fs = get_underlying_filesystem(path=from_path)
+    bucket, prefix = fs._split_path(from_path)  # pylint: disable=W0212
+    store: ObjectStore = fs._construct_store(bucket)
+
+    download_kwargs = {}
+    if "chunk_size" in kwargs:
+        download_kwargs["chunk_size"] = kwargs["chunk_size"]
+    if "max_concurrency" in kwargs:
+        download_kwargs["max_concurrency"] = kwargs["max_concurrency"]
+
+    reader = ObstoreParallelReader(store, **download_kwargs)
+    target_path = pathlib.Path(to_path) if isinstance(to_path, str) else to_path
+
+    # if recursive, just download the prefix to the target path
+    if recursive:
+        logger.debug(f"Downloading recursively {prefix=} to {target_path=}")
+        await reader.download_files(
+            prefix,
+            target_path,
+            exclude=exclude,
+        )
+        return str(to_path)
+
+    # if not recursive, we need to split out the file name from the prefix
+    else:
+        path_for_reader = pathlib.Path(prefix).name
+        final_prefix = pathlib.Path(prefix).parent
+        logger.debug(f"Downloading single file {final_prefix=}, {path_for_reader=} to {target_path=}")
+        await reader.download_files(
+            final_prefix,
+            target_path.parent,
+            path_for_reader,
+            destination_file_name=target_path.name,
+        )
+        return str(target_path)
+
+
 async def get(from_path: str, to_path: Optional[str | pathlib.Path] = None, recursive: bool = False, **kwargs) -> str:
     if not to_path:
-        name = pathlib.Path(from_path).name
+        name = pathlib.Path(from_path).name  # may need to be adjusted for windows
         to_path = get_random_local_path(file_path_or_file_name=name)
         logger.debug(f"Storing file from {from_path} to {to_path}")
+    else:
+        # Only apply directory logic for single files (not recursive)
+        if not recursive:
+            to_path_str = str(to_path)
+            # Check for trailing separator BEFORE converting to Path (which normalizes and removes it)
+            ends_with_sep = to_path_str.endswith(os.sep)
+            to_path_obj = pathlib.Path(to_path)
+
+            # If path ends with os.sep or is an existing directory, append source filename
+            if ends_with_sep or (to_path_obj.exists() and to_path_obj.is_dir()):
+                source_filename = pathlib.Path(from_path).name  # may need to be adjusted for windows
+                to_path = to_path_obj / source_filename
+        # For recursive=True, keep to_path as-is (it's the destination directory for contents)
+
     file_system = get_underlying_filesystem(path=from_path)
+
+    # Check if we should use obstore bypass
+    if (
+        _is_obstore_supported_protocol(file_system.protocol)
+        and hasattr(file_system, "_split_path")
+        and hasattr(file_system, "_construct_store")
+        and recursive
+    ):
+        return await _get_obstore_bypass(from_path, to_path, recursive, **kwargs)
+
     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}")
+        logger.debug(f"Error in getting {from_path} to {to_path}, recursive: {recursive}, error: {oe}")
         if isinstance(file_system, AsyncFileSystem):
             try:
                 exists = await file_system._exists(from_path)  # pylint: disable=W0212
@@ -103,7 +238,6 @@ async def get(from_path: str, to_path: Optional[str | pathlib.Path] = None, recu
         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}")
@@ -118,21 +252,21 @@ async def _get_from_filesystem(
     **kwargs,
 ):
     if isinstance(file_system, AsyncFileSystem):
-        dst = await file_system._get(from_path, to_path, recursive=recursive, **kwargs)  # pylint: disable=W0212
+        dst = await file_system._get(str(from_path), str(to_path), recursive=recursive, **kwargs)  # pylint: disable=W0212
     else:
-        dst = file_system.get(from_path, to_path, recursive=recursive, **kwargs)
+        dst = file_system.get(str(from_path), str(to_path), recursive=recursive, **kwargs)
 
     if isinstance(dst, (str, pathlib.Path)):
         return dst
-    return to_path
+    return str(to_path)
 
 
-async def put(from_path: str, to_path: Optional[str] = None, recursive: bool = False, **kwargs):
+async def put(from_path: str, to_path: Optional[str] = None, recursive: bool = False, **kwargs) -> str:
     if not to_path:
         from flyte._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
+        name = pathlib.Path(from_path).name
         to_path = ctx.raw_data.get_random_remote_path(file_name=name)
 
     file_system = get_underlying_filesystem(path=to_path)
@@ -142,11 +276,53 @@ async def put(from_path: str, to_path: Optional[str] = None, recursive: bool = F
     else:
         dst = file_system.put(from_path, to_path, recursive=recursive, **kwargs)
     if isinstance(dst, (str, pathlib.Path)):
-        return dst
+        return str(dst)
     else:
         return to_path
 
 
+async def _open_obstore_bypass(path: str, mode: str = "rb", **kwargs) -> AsyncReadableFile | AsyncWritableFile:
+    """
+    Simple obstore bypass for opening files. No fallbacks, obstore only.
+    """
+
+    fs = get_underlying_filesystem(path=path)
+    bucket, file_path = fs._split_path(path)  # pylint: disable=W0212
+    store: ObjectStore = fs._construct_store(bucket)
+
+    file_handle: AsyncReadableFile | AsyncWritableFile
+
+    if "w" in mode:
+        attributes = kwargs.pop("attributes", {})
+        file_handle = obstore.open_writer_async(store, file_path, attributes=attributes)
+    else:  # read mode
+        buffer_size = kwargs.pop("buffer_size", 10 * 2**20)
+        file_handle = await obstore.open_reader_async(store, file_path, buffer_size=buffer_size)
+    return file_handle
+
+
+async def open(path: str, mode: str = "rb", **kwargs) -> AsyncReadableFile | AsyncWritableFile:
+    """
+    Asynchronously open a file and return an async context manager.
+    This function checks if the underlying filesystem supports obstore bypass.
+    If it does, it uses obstore to open the file. Otherwise, it falls back to
+    the standard _open function which uses AsyncFileSystem.
+
+    It will raise NotImplementedError if neither obstore nor AsyncFileSystem is supported.
+    """
+    fs = get_underlying_filesystem(path=path)
+
+    # Check if we should use obstore bypass
+    if _is_obstore_supported_protocol(fs.protocol) and hasattr(fs, "_split_path") and hasattr(fs, "_construct_store"):
+        return await _open_obstore_bypass(path, mode, **kwargs)
+
+    # Fallback to normal open
+    if isinstance(fs, AsyncFileSystem):
+        return await fs.open_async(path, mode, **kwargs)
+
+    raise OnlyAsyncIOSupportedError(f"Filesystem {fs} does not support async operations")
+
+
 async def put_stream(
     data_iterable: typing.AsyncIterable[bytes] | bytes, *, name: str | None = None, to_path: str | None = None, **kwargs
 ) -> str:
@@ -172,70 +348,75 @@ async def put_stream(
 
         ctx = internal_ctx()
         to_path = ctx.raw_data.get_random_remote_path(file_name=name)
+
+    # Check if we should use obstore bypass
     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 str(to_path)
-        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:
+    try:
+        file_handle = typing.cast("AsyncWritableFile", await open(to_path, "wb", **kwargs))
         if isinstance(data_iterable, bytes):
-            f.write(data_iterable)
+            await file_handle.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)
+                await file_handle.write(data)
+        await file_handle.close()
+        return str(to_path)
+    except OnlyAsyncIOSupportedError:
+        pass
+
+    # Fallback to normal open
+    file_handle_io: typing.IO = fs.open(to_path, mode="wb", **kwargs)
+    if isinstance(data_iterable, bytes):
+        file_handle_io.write(data_iterable)
+    else:
+        async for data in data_iterable:
+            file_handle_io.write(data)
+    file_handle_io.close()
+
     return str(to_path)
 
 
-async def get_stream(path: str, chunk_size=10 * 2**20, **kwargs) -> AsyncIterator[bytes]:
+async def get_stream(path: str, chunk_size=10 * 2**20, **kwargs) -> AsyncGenerator[bytes, None]:
     """
     Get a stream of data from a remote location.
     This is useful for downloading streaming data from a remote location.
     Example usage:
     ```python
     import flyte.storage as storage
-    obj = storage.get_stream(path="s3://my_bucket/my_file.txt")
+    async for chunk in storage.get_stream(path="s3://my_bucket/my_file.txt"):
+        process(chunk)
     ```
 
     :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.
+    :return: An async iterator that yields chunks of bytes.
     """
-    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)
+    # Check if we should use obstore bypass
+    fs = get_underlying_filesystem(path=path)
+    if _is_obstore_supported_protocol(fs.protocol) and hasattr(fs, "_split_path") and hasattr(fs, "_construct_store"):
+        # Set buffer_size for obstore if chunk_size is provided
+        if "buffer_size" not in kwargs:
+            kwargs["buffer_size"] = chunk_size
+        file_handle = typing.cast("AsyncReadableFile", await _open_obstore_bypass(path, "rb", **kwargs))
+        while chunk := await file_handle.read():
+            yield bytes(chunk)
+        return
+
+    # Fallback to normal open
+    if "block_size" not in kwargs:
+        kwargs["block_size"] = chunk_size
+
+    if isinstance(fs, AsyncFileSystem):
+        file_handle = await fs.open_async(path, "rb", **kwargs)
+        while chunk := await file_handle.read():
             yield chunk
+        await file_handle.close()
+        return
+
+    file_handle = fs.open(path, "rb", **kwargs)
+    while chunk := file_handle.read():
+        yield chunk
+    file_handle.close()
 
 
 def join(*paths: str) -> str:
@@ -248,4 +429,32 @@ def join(*paths: str) -> str:
     return str(os.path.join(*paths))
 
 
-register(["s3", "gs", "abfs", "abfss"], asynchronous=True)
+async def exists(path: str, **kwargs) -> bool:
+    """
+    Check if a path exists.
+
+    :param path: Path to be checked.
+    :param kwargs: Additional arguments to be passed to the underlying filesystem.
+    :return: True if the path exists, False otherwise.
+    """
+    try:
+        fs = get_underlying_filesystem(path=path, **kwargs)
+        if isinstance(fs, AsyncFileSystem):
+            _ = await fs._info(path)
+            return True
+        _ = fs.info(path)
+        return True
+    except FileNotFoundError:
+        return False
+
+
+def exists_sync(path: str, **kwargs) -> bool:
+    try:
+        fs = get_underlying_filesystem(path=path, **kwargs)
+        _ = fs.info(path)
+        return True
+    except FileNotFoundError:
+        return False
+
+
+register(_OBSTORE_SUPPORTED_PROTOCOLS, asynchronous=True)

flyte/syncify/__init__.py

@@ -0,0 +1,56 @@
+"""
+# Syncify Module
+This module provides the `syncify` decorator and the `Syncify` class.
+The decorator can be used to convert asynchronous functions or methods into synchronous ones.
+This is useful for integrating async code into synchronous contexts.
+
+Every asynchronous function or method wrapped with `syncify` can be called synchronously using the
+parenthesis `()` operator, or asynchronously using the `.aio()` method.
+
+Example::
+
+```python
+from flyte.syncify import syncify
+
+@syncify
+async def async_function(x: str) -> str:
+    return f"Hello, Async World {x}!"
+
+
+# now you can call it synchronously
+result = async_function("Async World")   # Note: no .aio() needed for sync calls
+print(result)
+# Output: Hello, Async World Async World!
+
+# or call it asynchronously
+async def main():
+    result = await async_function.aio("World")    # Note the use of .aio() for async calls
+    print(result)
+```
+
+## Creating a Syncify Instance
+```python
+from flyte.syncify. import Syncify
+
+syncer = Syncify("my_syncer")
+
+# Now you can use `syncer` to decorate your async functions or methods
+
+```
+
+## How does it work?
+The Syncify class wraps asynchronous functions, classmethods, instance methods, and static methods to
+ provide a synchronous interface. The wrapped methods are always executed in the context of a background loop,
+ whether they are called synchronously or asynchronously. This allows for seamless integration of async code, as
+ certain async libraries capture the event loop. An example is grpc.aio, which captures the event loop.
+ In such a case, the Syncify class ensures that the async function is executed in the context of the background loop.
+
+To use it correctly with grpc.aio, you should wrap every grpc.aio channel creation, and client invocation
+with the same `Syncify` instance. This ensures that the async code runs in the correct event loop context.
+"""
+
+from flyte.syncify._api import Syncify
+
+syncify = Syncify()
+
+__all__ = ["Syncify", "syncify"]