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

flyte/storage/_parallel_reader.py

@@ -0,0 +1,274 @@
+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
+    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 = await inq.get()
+                    if task is sentinel:
+                        inq.put_nowait(sentinel)
+                        break
+                    chunk_source_offset = task.chunk.offset
+                    buf = active[task.source.id]
+                    data_to_write = await obstore.get_range_async(
+                        self._store,
+                        str(task.source.path),
+                        start=chunk_source_offset,
+                        end=chunk_source_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)
+                    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
+    ) -> 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
+        """
+
+        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_
+                    x = await obstore.head_async(self._store, str(path))
+                    yield x
+                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) -> 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

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"]
 
@@ -141,12 +147,76 @@ 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, **kwargs) -> str:
+    from obstore.store import ObjectStore
+
+    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,
+        )
+        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")
+    ):
+        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:
@@ -189,7 +259,7 @@ async def put(from_path: str, to_path: Optional[str] = None, recursive: bool = F
         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)
@@ -204,34 +274,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 +343,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 +377,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 +424,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/__init__.py

@@ -19,6 +19,10 @@ It is always possible to bypass the type system and use the `FlytePickle` type t
  written in python. The Pickled objects cannot be represented in the UI, and may be in-efficient for large datasets.
 """
 
+from importlib.metadata import entry_points
+
+from flyte._logging import logger
+
 from ._interface import guess_interface
 from ._pickle import FlytePickle
 from ._renderer import Renderable
@@ -34,3 +38,15 @@ __all__ = [
     "guess_interface",
     "literal_string_repr",
 ]
+
+
+def _load_custom_type_transformers():
+    plugins = entry_points(group="flyte.plugins.types")
+    for ep in plugins:
+        try:
+            logger.info(f"Loading type transformer: {ep.name}")
+            loaded = ep.load()
+            if callable(loaded):
+                loaded()
+        except Exception as e:
+            logger.warning(f"Failed to load type transformer {ep.name} with error: {e}")

flyte/types/_interface.py

@@ -1,9 +1,9 @@
 import inspect
 from typing import Any, Dict, Iterable, Tuple, Type, cast
 
-from flyteidl.core import interface_pb2, literals_pb2
+from flyteidl2.core import interface_pb2, literals_pb2
+from flyteidl2.task import common_pb2
 
-from flyte._protos.workflow import common_pb2
 from flyte.models import NativeInterface