flyte 2.0.0b22__py3-none-any.whl → 2.0.0b24__py3-none-any.whl

flyte/storage/_storage.py

@@ -1,3 +1,5 @@
+from __future__ import annotations
+
 import os
 import pathlib
 import random
@@ -7,6 +9,7 @@ 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
@@ -14,7 +17,10 @@ from obstore.fsspec import register
 
 from flyte._initialize import get_storage
 from flyte._logging import logger
-from flyte.errors import InitializationError
+from flyte.errors import InitializationError, OnlyAsyncIOSupportedError
+
+if typing.TYPE_CHECKING:
+    from obstore import AsyncReadableFile, AsyncWritableFile
 
 _OBSTORE_SUPPORTED_PROTOCOLS = ["s3", "gs", "abfs", "abfss"]
 
@@ -204,34 +210,48 @@ async def put(from_path: str, to_path: Optional[str] = None, recursive: bool = F
         return to_path
 
 
-async def _put_stream_obstore_bypass(data_iterable: typing.AsyncIterable[bytes] | bytes, to_path: str, **kwargs) -> str:
+async def _open_obstore_bypass(path: str, mode: str = "rb", **kwargs) -> AsyncReadableFile | AsyncWritableFile:
     """
-    NOTE: This can break if obstore changes its API.
-
-    This function is a workaround for obstore's fsspec implementation which does not support async file operations.
-    It uses the synchronous methods directly to put a stream of data.
+    Simple obstore bypass for opening files. No fallbacks, obstore only.
     """
-    import obstore
     from obstore.store import ObjectStore
 
-    fs = get_underlying_filesystem(path=to_path)
-    if not hasattr(fs, "_split_path") or not hasattr(fs, "_construct_store"):
-        raise NotImplementedError(f"Obstore bypass not supported for {fs.protocol} protocol, methods missing.")
-    bucket, path = fs._split_path(to_path)  # pylint: disable=W0212
+    fs = get_underlying_filesystem(path=path)
+    bucket, file_path = fs._split_path(path)  # pylint: disable=W0212
     store: ObjectStore = fs._construct_store(bucket)
-    if "attributes" in kwargs:
-        attributes = kwargs.pop("attributes")
-    else:
-        attributes = {}
-    buf_file = obstore.open_writer_async(store, path, attributes=attributes)
-    if isinstance(data_iterable, bytes):
-        await buf_file.write(data_iterable)
-    else:
-        async for data in data_iterable:
-            await buf_file.write(data)
-    # await buf_file.flush()
-    await buf_file.close()
-    return to_path
+
+    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(
@@ -259,60 +279,31 @@ async def put_stream(
 
         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:
-            if _is_obstore_supported_protocol(fs.protocol):
-                # If the protocol is supported by obstore, use the obstore bypass method
-                return await _put_stream_obstore_bypass(data_iterable, to_path=to_path, **kwargs)
-            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 as e:
-            logger.debug(f"{fs} doesn't implement 'open_async', falling back to sync, {e}")
-        finally:
-            if file_handle is not None:
-                await file_handle.close()
-
-    with fs.open(to_path, "wb", **kwargs) as f:
+    # Check if we should use obstore bypass
+    fs = get_underlying_filesystem(path=to_path)
+    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)
-    return str(to_path)
-
-
-async def _get_stream_obstore_bypass(path: str, chunk_size, **kwargs) -> AsyncGenerator[bytes, None]:
-    """
-    NOTE: This can break if obstore changes its API.
-    This function is a workaround for obstore's fsspec implementation which does not support async file operations.
-    It uses the synchronous methods directly to get a stream of data.
-    """
-    import obstore
-    from obstore.store import ObjectStore
+                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()
 
-    fs = get_underlying_filesystem(path=path)
-    if not hasattr(fs, "_split_path") or not hasattr(fs, "_construct_store"):
-        raise NotImplementedError(f"Obstore bypass not supported for {fs.protocol} protocol, methods missing.")
-    bucket, rem_path = fs._split_path(path)  # pylint: disable=W0212
-    store: ObjectStore = fs._construct_store(bucket)
-    buf_file = await obstore.open_reader_async(store, rem_path, buffer_size=chunk_size)
-    try:
-        while True:
-            chunk = await buf_file.read()
-            if not chunk:
-                break
-            yield bytes(chunk)
-    finally:
-        buf_file.close()
+    return str(to_path)
 
 
 async def get_stream(path: str, chunk_size=10 * 2**20, **kwargs) -> AsyncGenerator[bytes, None]:
@@ -322,42 +313,41 @@ async def get_stream(path: str, chunk_size=10 * 2**20, **kwargs) -> AsyncGenerat
     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)
+    # 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
 
-    file_size = fs.info(path)["size"]
-    total_read = 0
-    file_handle = None
-    try:
-        if _is_obstore_supported_protocol(fs.protocol):
-            # If the protocol is supported by obstore, use the obstore bypass method
-            async for x in _get_stream_obstore_bypass(path, chunk_size=chunk_size, **kwargs):
-                yield x
-            return
-        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 as e:
-        logger.debug(f"{fs} doesn't implement 'open_async', falling back to sync, error: {e}")
-    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)
+    # 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:
@@ -370,4 +360,32 @@ def join(*paths: str) -> str:
     return str(os.path.join(*paths))
 
 
+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/types/_pickle.py

@@ -116,14 +116,27 @@ class FlytePickleTransformer(TypeTransformer[FlytePickle]):
             and literal_type.blob.format == FlytePickleTransformer.PYTHON_PICKLE_FORMAT
         ):
             return FlytePickle
+        if literal_type.simple == types_pb2.SimpleType.BINARY:
+            return FlytePickle
 
         raise ValueError(f"Transformer {self} cannot reverse {literal_type}")
 
     def get_literal_type(self, t: Type[T]) -> types_pb2.LiteralType:
         lt = types_pb2.LiteralType(
-            blob=types_pb2.BlobType(
-                format=self.PYTHON_PICKLE_FORMAT, dimensionality=types_pb2.BlobType.BlobDimensionality.SINGLE
-            )
+            union_type=types_pb2.UnionType(
+                variants=[
+                    types_pb2.LiteralType(
+                        blob=types_pb2.BlobType(
+                            format=self.PYTHON_PICKLE_FORMAT,
+                            dimensionality=types_pb2.BlobType.BlobDimensionality.SINGLE,
+                        ),
+                        structure=types_pb2.TypeStructure(tag=self.name),
+                    ),
+                    types_pb2.LiteralType(
+                        simple=types_pb2.SimpleType.BINARY, structure=types_pb2.TypeStructure(tag=self.name)
+                    ),
+                ]
+            ),
         )
         lt.metadata = {"python_class_name": str(t)}
         return lt

{flyte-2.0.0b22.data → flyte-2.0.0b24.data}/scripts/runtime.py

@@ -12,6 +12,8 @@ from typing import Any, List
 
 import click
 
+from flyte.models import PathRewrite
+
 # Todo: work with pvditt to make these the names
 # ACTION_NAME = "_U_ACTION_NAME"
 # RUN_NAME = "_U_RUN_NAME"
@@ -28,8 +30,8 @@ ENDPOINT_OVERRIDE = "_U_EP_OVERRIDE"
 RUN_OUTPUT_BASE_DIR = "_U_RUN_BASE"
 FLYTE_ENABLE_VSCODE_KEY = "_F_E_VS"
 
-# TODO: Remove this after proper auth is implemented
 _UNION_EAGER_API_KEY_ENV_VAR = "_UNION_EAGER_API_KEY"
+_F_PATH_REWRITE = "_F_PATH_REWRITE"
 
 
 @click.group()
@@ -94,6 +96,7 @@ def main(
     import flyte
     import flyte._utils as utils
     import flyte.errors
+    import flyte.storage as storage
     from flyte._initialize import init
     from flyte._internal.controllers import create_controller
     from flyte._internal.imagebuild.image_builder import ImageCache
@@ -136,19 +139,34 @@ def main(
             controller_kwargs["insecure"] = True
         logger.debug(f"Using controller endpoint: {ep} with kwargs: {controller_kwargs}")
 
-    bundle = CodeBundle(tgz=tgz, pkl=pkl, destination=dest, computed_version=version)
+    bundle = None
+    if tgz or pkl:
+        bundle = CodeBundle(tgz=tgz, pkl=pkl, destination=dest, computed_version=version)
     init(org=org, project=project, domain=domain, image_builder="remote", **controller_kwargs)
     # Controller is created with the same kwargs as init, so that it can be used to run tasks
     controller = create_controller(ct="remote", **controller_kwargs)
 
     ic = ImageCache.from_transport(image_cache) if image_cache else None
 
+    path_rewrite_cfg = os.getenv(_F_PATH_REWRITE, None)
+    path_rewrite = None
+    if path_rewrite_cfg:
+        potential_path_rewrite = PathRewrite.from_str(path_rewrite_cfg)
+        if storage.exists_sync(potential_path_rewrite.new_prefix):
+            path_rewrite = potential_path_rewrite
+            logger.info(f"Path rewrite configured for {path_rewrite.new_prefix}")
+        else:
+            logger.error(
+                f"Path rewrite failed for path {potential_path_rewrite.new_prefix}, "
+                f"not found, reverting to original path {potential_path_rewrite.old_prefix}"
+            )
+
     # Create a coroutine to load the task and run it
     task_coroutine = load_and_run_task(
         resolver=resolver,
         resolver_args=resolver_args,
         action=ActionID(name=name, run_name=run_name, project=project, domain=domain, org=org),
-        raw_data_path=RawDataPath(path=raw_data_path),
+        raw_data_path=RawDataPath(path=raw_data_path, path_rewrite=path_rewrite),
         checkpoints=Checkpoints(checkpoint_path, prev_checkpoint),
         code_bundle=bundle,
         input_path=inputs,
@@ -166,8 +184,20 @@ def main(
     async def _run_and_stop():
         loop = asyncio.get_event_loop()
         loop.set_exception_handler(flyte.errors.silence_grpc_polling_error)
-        await utils.run_coros(controller_failure, task_coroutine)
-        await controller.stop()
+        try:
+            await utils.run_coros(controller_failure, task_coroutine)
+            await controller.stop()
+        except flyte.errors.RuntimeSystemError as e:
+            logger.error(f"Runtime system error: {e}")
+            from flyte._internal.runtime.convert import convert_from_native_to_error
+            from flyte._internal.runtime.io import upload_error
+
+            logger.error(f"Flyte runtime failed for action {name} with run name {run_name}, error: {e}")
+            err = convert_from_native_to_error(e)
+            path = await upload_error(err.err, outputs_path)
+            logger.error(f"Run {run_name} Action {name} failed with error: {err}. Uploaded error to {path}")
+            await controller.stop()
+            raise
 
     asyncio.run(_run_and_stop())
     logger.warning(f"Flyte runtime completed for action {name} with run name {run_name}")

{flyte-2.0.0b22.dist-info → flyte-2.0.0b24.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: flyte
-Version: 2.0.0b22
+Version: 2.0.0b24
 Summary: Add your description here
 Author-email: Ketan Umare <kumare3@users.noreply.github.com>
 Requires-Python: >=3.10
@@ -25,6 +25,8 @@ Requires-Dist: async-lru>=2.0.5
 Requires-Dist: mashumaro
 Requires-Dist: dataclasses_json
 Requires-Dist: aiolimiter>=1.2.1
+Provides-Extra: aiosqlite
+Requires-Dist: aiosqlite>=0.21.0; extra == "aiosqlite"
 Dynamic: license-file
 
 # Flyte 2 SDK 🚀