flyte 2.0.0b32__py3-none-any.whl

flyte/syncify/_api.py

@@ -0,0 +1,375 @@
+from __future__ import annotations
+
+import asyncio
+import atexit
+import concurrent.futures
+import functools
+import inspect
+import logging
+import threading
+from typing import (
+    Any,
+    AsyncIterator,
+    Awaitable,
+    Callable,
+    Coroutine,
+    Iterator,
+    ParamSpec,
+    Protocol,
+    TypeVar,
+    Union,
+    cast,
+    overload,
+)
+
+from flyte._logging import logger
+
+P = ParamSpec("P")
+R_co = TypeVar("R_co", covariant=True)
+T = TypeVar("T")
+
+
+class SyncFunction(Protocol[P, R_co]):
+    """
+    A protocol that defines the interface for synchronous functions or methods that can be converted from asynchronous
+     ones.
+    """
+
+    def __call__(self, *args: Any, **kwargs: Any) -> R_co: ...
+
+    def aio(self, *args: Any, **kwargs: Any) -> Awaitable[R_co]: ...
+
+
+class SyncGenFunction(Protocol[P, R_co]):
+    """
+    A protocol that defines the interface for synchronous functions or methods that can be converted from asynchronous
+     ones.
+    """
+
+    def __call__(self, *args: Any, **kwargs: Any) -> Iterator[R_co]: ...
+
+    def aio(self, *args: Any, **kwargs: Any) -> AsyncIterator[R_co]: ...
+
+
+class _BackgroundLoop:
+    """
+    A background event loop that runs in a separate thread and used the `Syncify` decorator to run asynchronous
+    functions or methods synchronously.
+    """
+
+    def __init__(self, name: str):
+        self.loop = asyncio.new_event_loop()
+        self.thread = threading.Thread(name=name, target=self._run, daemon=True)
+        self.thread.start()
+        atexit.register(self.stop)
+
+    def _run(self):
+        import flyte.errors
+
+        # Set the exception handler to silence specific gRPC polling errors
+        self.loop.set_exception_handler(flyte.errors.silence_grpc_polling_error)
+        asyncio.set_event_loop(self.loop)
+        self.loop.run_forever()
+
+    def stop(self):
+        # stop the loop and wait briefly for thread to exit
+        self.loop.call_soon_threadsafe(self.loop.stop)
+        self.thread.join(timeout=1)
+
+    def is_in_loop(self) -> bool:
+        """
+        Check if the current thread is the background loop thread.
+        """
+        # If the current thread is not the background loop thread, return False
+        if threading.current_thread() != self.thread:
+            return False
+
+        if not self.thread.is_alive():
+            # If the thread is not alive, we cannot be in the loop
+            return False
+
+        # Lets get the current event loop and check if it matches the background loop
+        loop = None
+        try:
+            loop = asyncio.get_running_loop()
+        except RuntimeError:
+            pass
+
+        return loop == self.loop
+
+    def iterate_in_loop_sync(self, async_gen: AsyncIterator[R_co]) -> Iterator[R_co]:
+        # Create an iterator that pulls items from the async generator
+        assert self.thread.name != threading.current_thread().name, (
+            f"Cannot run coroutine in the same thread {self.thread.name}"
+        )
+        while True:
+            try:
+                # use __anext__() and cast to Coroutine so mypy is happy
+                future: concurrent.futures.Future[R_co] = asyncio.run_coroutine_threadsafe(
+                    cast(Coroutine[Any, Any, R_co], async_gen.__anext__()),
+                    self.loop,
+                )
+                yield future.result()
+            except (StopAsyncIteration, StopIteration):
+                break
+            except Exception as e:
+                if logger.getEffectiveLevel() > logging.DEBUG:
+                    # If the log level is not DEBUG, we will remove the extra stack frames to avoid confusion for the
+                    # user
+                    # This is because the stack trace will include the Syncify wrapper and the background loop thread
+                    tb = e.__traceback__
+                    while tb and tb.tb_next:
+                        if tb.tb_frame.f_code.co_name == "":
+                            break
+                        tb = tb.tb_next
+                    raise e.with_traceback(tb)
+                # If the log level is DEBUG, we will keep the extra stack frames to help with debugging
+                raise e
+
+    def call_in_loop_sync(self, coro: Coroutine[Any, Any, R_co]) -> R_co | Iterator[R_co]:
+        """
+        Run the given coroutine in the background loop and return its result.
+        """
+        future: concurrent.futures.Future[R_co | AsyncIterator[R_co]] = asyncio.run_coroutine_threadsafe(
+            coro, self.loop
+        )
+        result = future.result()
+        if result is not None and hasattr(result, "__aiter__"):
+            # If the result is an async iterator, we need to convert it to a sync iterator
+            return cast(Iterator[R_co], self.iterate_in_loop_sync(cast(AsyncIterator[R_co], result)))
+        # Otherwise, just return the result
+        return result
+
+    async def iterate_in_loop(self, async_gen: AsyncIterator[R_co]) -> AsyncIterator[R_co]:
+        """
+        Run the given async iterator in the background loop and yield its results.
+        """
+        if self.is_in_loop():
+            # If we are already in the background loop, just return the async iterator
+            async for r in async_gen:
+                yield r
+            return
+
+        while True:
+            try:
+                # same replacement here for the async path
+                future: concurrent.futures.Future[R_co] = asyncio.run_coroutine_threadsafe(
+                    cast(Coroutine[Any, Any, R_co], async_gen.__anext__()),
+                    self.loop,
+                )
+                # Wrap the future in an asyncio Future to yield it in an async context
+                aio_future: asyncio.Future[R_co] = asyncio.wrap_future(future)
+                # await for the future to complete and yield its result
+                v = await aio_future
+                yield v
+            except StopAsyncIteration:
+                break
+            except Exception as e:
+                if logger.getEffectiveLevel() > logging.DEBUG:
+                    # If the log level is not DEBUG, we will remove the extra stack frames to avoid confusion for the
+                    # user.
+                    # This is because the stack trace will include the Syncify wrapper and the background loop thread
+                    tb = e.__traceback__
+                    while tb and tb.tb_next:
+                        if tb.tb_frame.f_code.co_name == "":
+                            break
+                        tb = tb.tb_next
+                    raise e.with_traceback(tb)
+                # If the log level is DEBUG, we will keep the extra stack frames to help with debugging
+                raise e
+
+    async def aio(self, coro: Coroutine[Any, Any, R_co]) -> R_co:
+        """
+        Run the given coroutine in the background loop and return its result.
+        """
+        if self.is_in_loop():
+            # If we are already in the background loop, just run the coroutine
+            return await coro
+        try:
+            # Otherwise, run it in the background loop and wait for the result
+            future: concurrent.futures.Future[R_co] = asyncio.run_coroutine_threadsafe(coro, self.loop)
+            # Wrap the future in an asyncio Future to await it in an async context
+            aio_future: asyncio.Future[R_co] = asyncio.wrap_future(future)
+            # await for the future to complete and return its result
+            return await aio_future
+        except Exception as e:
+            if logger.getEffectiveLevel() > logging.DEBUG:
+                # If the log level is not DEBUG, we will remove the extra stack frames to avoid confusion for the user
+                # This is because the stack trace will include the Syncify wrapper and the background loop thread
+                tb = e.__traceback__
+                while tb and tb.tb_next:
+                    if tb.tb_frame.f_code.co_name == "":
+                        break
+                    tb = tb.tb_next
+                raise e.with_traceback(tb)
+            # If the log level is DEBUG, we will keep the extra stack frames to help with debugging
+            raise e
+
+
+class _SyncWrapper:
+    """
+    A wrapper class that the Syncify decorator uses to convert asynchronous functions or methods into synchronous ones.
+    """
+
+    def __init__(
+        self,
+        fn: Any,
+        bg_loop: _BackgroundLoop,
+        underlying_obj: Any = None,
+    ):
+        self.fn = fn
+        self._bg_loop = bg_loop
+        self._underlying_obj = underlying_obj
+
+    def __get__(self, instance: Any, owner: Any) -> Any:
+        """
+        This method is called when the wrapper is accessed as a method of a class instance.
+        :param instance:
+        :param owner:
+        :return:
+        """
+        fn: Any = self.fn
+        if instance is not None:
+            # If we have an instance, we need to bind the method to the instance (for instance methods)
+            fn = self.fn.__get__(instance, owner)
+
+        if instance is None and owner is not None and self._underlying_obj is not None:
+            # If we have an owner, we need to bind the method to the owner (for classmethods or staticmethods)
+            fn = self._underlying_obj.__get__(None, owner)
+
+        wrapper = _SyncWrapper(fn, bg_loop=self._bg_loop, underlying_obj=self._underlying_obj)
+        functools.update_wrapper(wrapper, self.fn)
+        return wrapper
+
+    def __call__(self, *args: Any, **kwargs: Any) -> Any:
+        if threading.current_thread().name == self._bg_loop.thread.name:
+            # If we are already in the background loop thread, we can call the function directly
+            raise AssertionError(
+                f"Deadlock detected: blocking call used in syncify thread {self._bg_loop.thread.name} "
+                f"when calling function {self.fn}, use .aio() if in an async call."
+            )
+        try:
+            # bind method if needed
+            coro_fn = self.fn
+
+            if inspect.isasyncgenfunction(coro_fn):
+                # Handle async iterator by converting to sync iterator
+                async_gen = coro_fn(*args, **kwargs)
+                return self._bg_loop.iterate_in_loop_sync(async_gen)
+            else:
+                return self._bg_loop.call_in_loop_sync(coro_fn(*args, **kwargs))
+        except Exception as e:
+            if logger.getEffectiveLevel() > logging.DEBUG:
+                # If the log level is not DEBUG, we will remove the extra stack frames to avoid confusion for the user
+                # This is because the stack trace will include the Syncify wrapper and the background loop thread
+                tb = e.__traceback__
+                while tb and tb.tb_next:
+                    if tb.tb_frame.f_code.co_name == self.fn.__name__:
+                        break
+                    tb = tb.tb_next
+                raise e.with_traceback(tb)
+            # If the log level is DEBUG, we will keep the extra stack frames to help with debugging
+            raise e
+
+    def aio(self, *args: Any, **kwargs: Any) -> Any:
+        fn = self.fn
+
+        try:
+            if inspect.isasyncgenfunction(fn):
+                # If the function is an async generator, we need to handle it differently
+                async_iter = fn(*args, **kwargs)
+                return self._bg_loop.iterate_in_loop(async_iter)
+            else:
+                # If we are already in the background loop, just return the coroutine
+                coro = fn(*args, **kwargs)
+                if hasattr(coro, "__aiter__"):
+                    # If the coroutine is an async iterator, we need to handle it differently
+                    return self._bg_loop.iterate_in_loop(coro)
+                return self._bg_loop.aio(coro)
+        except Exception as e:
+            if logger.getEffectiveLevel() > logging.DEBUG:
+                # If the log level is not DEBUG, we will remove the extra stack frames to avoid confusion for the user
+                # This is because the stack trace will include the Syncify wrapper and the background loop thread
+                tb = e.__traceback__
+                while tb and tb.tb_next:
+                    if tb.tb_frame.f_code.co_name == self.fn.__name__:
+                        break
+                    tb = tb.tb_next
+                raise e.with_traceback(tb)
+            # If the log level is DEBUG, we will keep the extra stack frames to help with debugging
+            raise e
+
+
+class Syncify:
+    """
+    A decorator to convert asynchronous functions or methods into synchronous ones.
+
+    This is useful for integrating async code into synchronous contexts.
+
+    Example::
+
+    ```python
+    syncer = Syncify()
+
+    @syncer
+    async def async_function(x: str) -> str:
+        return f"Hello, Async World {x}!"
+
+
+    # now you can call it synchronously
+    result = async_function("Async World")
+    print(result)
+    # Output: Hello, Async World Async World!
+
+    # or call it asynchronously
+    async def main():
+        result = await async_function.aio("World")
+        print(result)
+    ```
+
+    """
+
+    def __init__(self, name: str = "flyte_syncify"):
+        self._bg_loop = _BackgroundLoop(name=name)
+
+    @overload
+    def __call__(self, func: Callable[P, Awaitable[R_co]]) -> Any: ...
+
+    # def __call__(self, func: Callable[P, Awaitable[R_co]]) -> SyncFunction[P, R_co]: ...
+
+    @overload
+    def __call__(self, func: Callable[P, Iterator[R_co] | AsyncIterator[R_co]]) -> SyncGenFunction[P, R_co]: ...
+
+    # def __call__(self, func: Callable[[Type[T], *P.args, *P.kwargs], Awaitable[R_co]])
+    # -> SyncFunction[[Type[T], *P.args, *P.kwargs], R_co]: ...
+    @overload
+    def __call__(self, func: classmethod) -> Union[SyncFunction[P, R_co], SyncGenFunction[P, R_co]]: ...
+
+    @overload
+    def __call__(self, func: staticmethod) -> staticmethod: ...
+
+    def __call__(self, obj):
+        if isinstance(obj, classmethod):
+            wrapper = _SyncWrapper(obj.__func__, bg_loop=self._bg_loop, underlying_obj=obj)
+            functools.update_wrapper(wrapper, obj.__func__)
+            return wrapper
+
+        if isinstance(obj, staticmethod):
+            fn = obj.__func__
+            wrapper = _SyncWrapper(fn, bg_loop=self._bg_loop)
+            functools.update_wrapper(wrapper, fn)
+            return staticmethod(wrapper)
+
+        if inspect.isasyncgenfunction(obj):
+            wrapper = _SyncWrapper(obj, bg_loop=self._bg_loop)
+            functools.update_wrapper(wrapper, obj)
+            return cast(Callable[P, Iterator[R_co]], wrapper)
+
+        if inspect.iscoroutinefunction(obj):
+            wrapper = _SyncWrapper(obj, bg_loop=self._bg_loop)
+            functools.update_wrapper(wrapper, obj)
+            return wrapper
+
+        raise TypeError(
+            "Syncify can only be applied to async functions, async generators, async classmethods or staticmethods."
+        )

flyte/types/__init__.py

@@ -0,0 +1,52 @@
+"""
+# Flyte Type System
+
+The Flyte type system provides a way to define, transform, and manipulate types in Flyte workflows.
+Since the data flowing through Flyte has to often cross process, container and langauge boundaries, the type system
+is designed to be serializable to a universal format that can be understood across different environments. This
+universal format is based on Protocol Buffers. The types are called LiteralTypes and the runtime
+representation of data is called Literals.
+
+The type system includes:
+- **TypeEngine**: The core engine that manages type transformations and serialization. This is the main entry point for
+  for all the internal type transformations and serialization logic.
+- **TypeTransformer**: A class that defines how to transform one type to another. This is extensible
+    allowing users to define custom types and transformations.
+- **Renderable**: An interface for types that can be rendered as HTML, that can be outputted to a flyte.report.
+
+It is always possible to bypass the type system and use the `FlytePickle` type to serialize any python object
+ into a pickle format. The pickle format is not human-readable, but can be passed between flyte tasks that are
+ 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
+from ._string_literals import literal_string_repr
+from ._type_engine import TypeEngine, TypeTransformer, TypeTransformerFailedError
+
+__all__ = [
+    "FlytePickle",
+    "Renderable",
+    "TypeEngine",
+    "TypeTransformer",
+    "TypeTransformerFailedError",
+    "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

@@ -0,0 +1,40 @@
+import inspect
+from typing import Any, Dict, Iterable, Tuple, Type, cast
+
+from flyteidl2.core import interface_pb2, literals_pb2
+from flyteidl2.task import common_pb2
+
+from flyte.models import NativeInterface
+
+
+def guess_interface(
+    interface: interface_pb2.TypedInterface, default_inputs: Iterable[common_pb2.NamedParameter] | None = None
+) -> NativeInterface:
+    """
+    Returns the interface of the task with guessed types, as types may not be present in current env.
+    """
+    import flyte.types
+
+    if interface is None:
+        return NativeInterface({}, {})
+
+    default_input_literals: Dict[str, literals_pb2.Literal] = {}
+    if default_inputs is not None:
+        for param in default_inputs:
+            if param.parameter.HasField("default"):
+                default_input_literals[param.name] = param.parameter.default
+
+    guessed_inputs: Dict[str, Tuple[Type[Any], Any] | Any] = {}
+    if interface.inputs is not None and len(interface.inputs.variables) > 0:
+        input_types = flyte.types.TypeEngine.guess_python_types(cast(dict, interface.inputs.variables))
+        for name, t in input_types.items():
+            if name not in default_input_literals:
+                guessed_inputs[name] = (t, inspect.Parameter.empty)
+            else:
+                guessed_inputs[name] = (t, NativeInterface.has_default)
+
+    guessed_outputs: Dict[str, Type[Any]] = {}
+    if interface.outputs is not None and len(interface.outputs.variables) > 0:
+        guessed_outputs = flyte.types.TypeEngine.guess_python_types(cast(dict, interface.outputs.variables))
+
+    return NativeInterface.from_types(guessed_inputs, guessed_outputs, default_input_literals)

flyte/types/_pickle.py

@@ -0,0 +1,145 @@
+import hashlib
+import os
+import sys
+import typing
+from typing import Type
+
+import aiofiles
+import cloudpickle
+from flyteidl2.core import literals_pb2, types_pb2
+
+import flyte.storage as storage
+
+from ._type_engine import TypeEngine, TypeTransformer
+
+T = typing.TypeVar("T")
+DEFAULT_PICKLE_BYTES_LIMIT = 2**10 * 10  # 10KB
+
+
+class FlytePickle(typing.Generic[T]):
+    """
+    This type is only used by flytekit internally. User should not use this type.
+    Any type that flyte can't recognize will become FlytePickle
+    """
+
+    @classmethod
+    def python_type(cls) -> typing.Type:
+        return type(None)
+
+    @classmethod
+    def __class_getitem__(cls, python_type: typing.Type) -> typing.Type:
+        if python_type is None:
+            return cls
+
+        class _SpecificFormatClass(FlytePickle):
+            # Get the type engine to see this as kind of a generic
+            __origin__ = FlytePickle
+
+            @classmethod
+            def python_type(cls) -> typing.Type:
+                return python_type
+
+        return _SpecificFormatClass
+
+    @classmethod
+    async def to_pickle(cls, python_val: typing.Any) -> str:
+        h = hashlib.md5()
+        str_bytes = cloudpickle.dumps(python_val)
+        h.update(str_bytes)
+
+        uri = storage.get_random_local_path(file_path_or_file_name=h.hexdigest())
+        os.makedirs(os.path.dirname(uri), exist_ok=True)
+        async with aiofiles.open(uri, "w+b") as outfile:
+            await outfile.write(str_bytes)
+
+        return await storage.put(str(uri))
+
+    @classmethod
+    async def from_pickle(cls, uri: str) -> typing.Any:
+        # Deserialize the pickle, and return data in the pickle,
+        # and download pickle file to local first if file is not in the local file systems.
+        if storage.is_remote(uri):
+            local_path = storage.get_random_local_path()
+            await storage.get(uri, str(local_path), False)
+            uri = str(local_path)
+        async with aiofiles.open(uri, "rb") as infile:
+            data = cloudpickle.loads(await infile.read())
+        return data
+
+
+class FlytePickleTransformer(TypeTransformer[FlytePickle]):
+    PYTHON_PICKLE_FORMAT = "PythonPickle"
+
+    def __init__(self):
+        super().__init__(name="FlytePickle", t=FlytePickle)
+
+    def assert_type(self, t: Type[T], v: T):
+        # Every type can serialize to pickle, so we don't need to check the type here.
+        ...
+
+    async def to_python_value(self, lv: literals_pb2.Literal, expected_python_type: Type[T]) -> T:
+        if lv.scalar.blob and lv.scalar.blob.uri:
+            uri = lv.scalar.blob.uri
+            return await FlytePickle.from_pickle(uri)
+        elif lv.scalar.binary and lv.scalar.binary.value:
+            return cloudpickle.loads(lv.scalar.binary.value)
+        else:
+            raise ValueError(f"Cannot convert {lv} to {expected_python_type}")
+
+    async def to_literal(
+        self,
+        python_val: T,
+        python_type: Type[T],
+        expected: types_pb2.LiteralType,
+    ) -> literals_pb2.Literal:
+        if python_val is None:
+            raise AssertionError("Cannot pickle None Value.")
+        meta = literals_pb2.BlobMetadata(
+            type=types_pb2.BlobType(
+                format=self.PYTHON_PICKLE_FORMAT, dimensionality=types_pb2.BlobType.BlobDimensionality.SINGLE
+            )
+        )
+        if sys.getsizeof(python_val) > DEFAULT_PICKLE_BYTES_LIMIT:
+            remote_path = await FlytePickle.to_pickle(python_val)
+            return literals_pb2.Literal(
+                scalar=literals_pb2.Scalar(blob=literals_pb2.Blob(metadata=meta, uri=remote_path))
+            )
+        else:
+            return literals_pb2.Literal(
+                scalar=literals_pb2.Scalar(binary=literals_pb2.Binary(value=cloudpickle.dumps(python_val)))
+            )
+
+    def guess_python_type(self, literal_type: types_pb2.LiteralType) -> typing.Type[FlytePickle[typing.Any]]:
+        if (
+            literal_type.blob is not None
+            and literal_type.blob.dimensionality == types_pb2.BlobType.BlobDimensionality.SINGLE
+            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(
+            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
+
+
+TypeEngine.register(FlytePickleTransformer())