flyte 0.2.0b8__py3-none-any.whl → 0.2.0b9__py3-none-any.whl

flyte/__init__.py

@@ -39,6 +39,7 @@ __all__ = [
     "group",
     "init",
     "init_auto_from_config",
+    "map",
     "run",
     "trace",
     "with_runcontext",
@@ -51,6 +52,7 @@ from ._environment import Environment
 from ._group import group
 from ._image import Image
 from ._initialize import init, init_auto_from_config
+from ._map import map
 from ._resources import GPU, TPU, Device, Resources
 from ._retry import RetryStrategy
 from ._reusable_environment import ReusePolicy

flyte/_context.py

@@ -4,6 +4,7 @@ import contextvars
 from dataclasses import dataclass, replace
 from typing import TYPE_CHECKING, Awaitable, Callable, Optional, ParamSpec, TypeVar
 
+from flyte._logging import logger
 from flyte.models import GroupData, RawDataPath, TaskContext
 
 if TYPE_CHECKING:
@@ -49,6 +50,7 @@ class Context:
             raise ValueError("Cannot create a new context without contextdata.")
         self._data = data
         self._id = id(self)  # Immutable unique identifier
+        self._token = None  # Context variable token to restore the previous context
 
     @property
     def data(self) -> ContextData:
@@ -106,7 +108,11 @@ class Context:
 
     def __exit__(self, exc_type, exc_val, exc_tb):
         """Exit the context, restoring the previous context."""
-        root_context_var.reset(self._token)
+        try:
+            root_context_var.reset(self._token)
+        except Exception as e:
+            logger.warn(f"Failed to reset context: {e}")
+            raise e
 
     async def __aenter__(self):
         """Async version of context entry."""

flyte/_group.py

@@ -29,3 +29,4 @@ def group(name: str):
     new_tctx = tctx.replace(group_data=GroupData(name))
     with ctx.replace_task_context(new_tctx):
         yield
+        # Exit the context and restore the previous context

flyte/_internal/controllers/__init__.py

@@ -1,5 +1,6 @@
+import concurrent.futures
 import threading
-from typing import Any, Callable, Literal, Optional, Protocol, Tuple, TypeVar
+from typing import TYPE_CHECKING, Any, Callable, Literal, Optional, Protocol, Tuple, TypeVar
 
 from flyte._task import TaskTemplate
 from flyte.models import ActionID, NativeInterface
@@ -10,6 +11,9 @@ __all__ = ["Controller", "ControllerType", "TraceInfo", "create_controller", "ge
 
 from ..._protos.workflow import task_definition_pb2
 
+if TYPE_CHECKING:
+    import concurrent.futures
+
 ControllerType = Literal["local", "remote"]
 
 R = TypeVar("R")
@@ -28,7 +32,14 @@ class Controller(Protocol):
         """
         ...
 
-    def submit_sync(self, _task: TaskTemplate, *args, **kwargs) -> Any: ...
+    def submit_sync(self, _task: TaskTemplate, *args, **kwargs) -> concurrent.futures.Future:
+        """
+        This should call the async submit method above, but return a concurrent Future object that can be
+        used in a blocking wait or wrapped in an async future. This is called when
+          a) a synchronous task is kicked off locally,
+          b) a running task (of either kind) kicks off a downstream synchronous task.
+        """
+        ...
 
     async def submit_task_ref(self, _task: task_definition_pb2.TaskDetails, *args, **kwargs) -> Any:
         """

flyte/_internal/controllers/_local_controller.py

@@ -1,3 +1,8 @@
+import asyncio
+import atexit
+import concurrent.futures
+import os
+import threading
 from typing import Any, Callable, Tuple, TypeVar
 
 import flyte.errors
@@ -8,20 +13,67 @@ from flyte._internal.runtime.entrypoints import direct_dispatch
 from flyte._logging import log, logger
 from flyte._protos.workflow import task_definition_pb2
 from flyte._task import TaskTemplate
-from flyte._utils.asyn import loop_manager
+from flyte._utils.helpers import _selector_policy
 from flyte.models import ActionID, NativeInterface, RawDataPath
 
 R = TypeVar("R")
 
 
+class _TaskRunner:
+    """A task runner that runs an asyncio event loop on a background thread."""
+
+    def __init__(self) -> None:
+        self.__loop: asyncio.AbstractEventLoop | None = None
+        self.__runner_thread: threading.Thread | None = None
+        self.__lock = threading.Lock()
+        atexit.register(self._close)
+
+    def _close(self) -> None:
+        if self.__loop:
+            self.__loop.stop()
+
+    def _execute(self) -> None:
+        loop = self.__loop
+        assert loop is not None
+        try:
+            loop.run_forever()
+        finally:
+            loop.close()
+
+    def get_exc_handler(self):
+        def exc_handler(loop, context):
+            logger.error(
+                f"Taskrunner for {self.__runner_thread.name if self.__runner_thread else 'no thread'} caught"
+                f" exception in {loop}: {context}"
+            )
+
+        return exc_handler
+
+    def get_run_future(self, coro: Any) -> concurrent.futures.Future:
+        """Synchronously run a coroutine on a background thread."""
+        name = f"{threading.current_thread().name} : loop-runner"
+        with self.__lock:
+            if self.__loop is None:
+                with _selector_policy():
+                    self.__loop = asyncio.new_event_loop()
+
+                exc_handler = self.get_exc_handler()
+                self.__loop.set_exception_handler(exc_handler)
+                self.__runner_thread = threading.Thread(target=self._execute, daemon=True, name=name)
+                self.__runner_thread.start()
+        fut = asyncio.run_coroutine_threadsafe(coro, self.__loop)
+        return fut
+
+
 class LocalController:
     def __init__(self):
         logger.debug("LocalController init")
+        self._runner_map: dict[str, _TaskRunner] = {}
 
     @log
     async def submit(self, _task: TaskTemplate, *args, **kwargs) -> Any:
         """
-        Submit a node to the controller
+        Main entrypoint for submitting a task to the local controller.
         """
         ctx = internal_ctx()
         tctx = ctx.data.task_context
@@ -59,7 +111,17 @@ class LocalController:
             return result
         return out
 
-    submit_sync = loop_manager.synced(submit)
+    def submit_sync(self, _task: TaskTemplate, *args, **kwargs) -> concurrent.futures.Future:
+        name = threading.current_thread().name + f"PID:{os.getpid()}"
+        coro = self.submit(_task, *args, **kwargs)
+        if name not in self._runner_map:
+            if len(self._runner_map) > 100:
+                logger.warning(
+                    "More than 100 event loop runners created!!! This could be a case of runaway recursion..."
+                )
+            self._runner_map[name] = _TaskRunner()
+
+        return self._runner_map[name].get_run_future(coro)
 
     async def finalize_parent_action(self, action: ActionID):
         pass

flyte/_internal/controllers/remote/_controller.py

@@ -1,6 +1,9 @@
 from __future__ import annotations
 
 import asyncio
+import concurrent.futures
+import os
+import threading
 from collections import defaultdict
 from collections.abc import Callable
 from pathlib import Path
@@ -21,7 +24,7 @@ from flyte._internal.runtime.task_serde import translate_task_to_wire
 from flyte._logging import logger
 from flyte._protos.workflow import run_definition_pb2, task_definition_pb2
 from flyte._task import TaskTemplate
-from flyte._utils.asyn import loop_manager
+from flyte._utils.helpers import _selector_policy
 from flyte.models import ActionID, NativeInterface, SerializationContext
 
 R = TypeVar("R")
@@ -120,6 +123,8 @@ class RemoteController(Controller):
         self._parent_action_task_call_sequence: DefaultDict[str, DefaultDict[int, int]] = defaultdict(
             lambda: defaultdict(int)
         )
+        self._submit_loop: asyncio.AbstractEventLoop | None = None
+        self._submit_thread: threading.Thread | None = None
 
     def generate_task_call_sequence(self, task_obj: object, action_id: ActionID) -> int:
         """
@@ -236,7 +241,47 @@ class RemoteController(Controller):
         async with self._parent_action_semaphore[unique_action_name(current_action_id)]:
             return await self._submit(task_call_seq, _task, *args, **kwargs)
 
-    submit_sync = loop_manager.synced(submit)
+    def _sync_thread_loop_runner(self) -> None:
+        """This method runs the event loop and should be invoked in a separate thread."""
+
+        loop = self._submit_loop
+        assert loop is not None
+        try:
+            loop.run_forever()
+        finally:
+            loop.close()
+
+    def submit_sync(self, _task: TaskTemplate, *args, **kwargs) -> concurrent.futures.Future:
+        """
+        This function creates a cached thread and loop for the purpose of calling the submit method synchronously,
+        returning a concurrent Future that can be awaited. There's no need for a lock because this function itself is
+        single threaded and non-async. This pattern here is basically the trivial/degenerate case of the thread pool
+        in the LocalController.
+        Please see additional comments in protocol.
+
+        :param _task:
+        :param args:
+        :param kwargs:
+        :return:
+        """
+        if self._submit_thread is None:
+            # Please see LocalController for the general implementation of this pattern.
+            def exc_handler(loop, context):
+                logger.error(f"Remote controller submit sync loop caught exception in {loop}: {context}")
+
+            with _selector_policy():
+                self._submit_loop = asyncio.new_event_loop()
+                self._submit_loop.set_exception_handler(exc_handler)
+
+            self._submit_thread = threading.Thread(
+                name=f"remote-controller-{os.getpid()}-submitter", daemon=True, target=self._sync_thread_loop_runner
+            )
+            self._submit_thread.start()
+
+        coro = self.submit(_task, *args, **kwargs)
+        assert self._submit_loop is not None, "Submit loop should always have been initialized by now"
+        fut = asyncio.run_coroutine_threadsafe(coro, self._submit_loop)
+        return fut
 
     async def finalize_parent_action(self, action_id: ActionID):
         """

flyte/_internal/runtime/taskrunner.py

@@ -136,8 +136,9 @@ async def convert_and_run(
         raw_data_path=raw_data_path,
         compiled_image_cache=image_cache,
         report=flyte.report.Report(name=action.name),
+        mode="remote" if not ctx.data.task_context else ctx.data.task_context.mode,
     )
-    async with ctx.replace_task_context(tctx):
+    with ctx.replace_task_context(tctx):
         inputs_kwargs = await convert_inputs_to_native(inputs, task.native_interface)
         out, err = await run_task(tctx=tctx, controller=controller, task=task, inputs=inputs_kwargs)
         if err is not None:

flyte/_map.py

@@ -0,0 +1,215 @@
+import asyncio
+from typing import Any, AsyncGenerator, AsyncIterator, Generic, Iterable, Iterator, List, Union, cast
+
+from flyte.syncify import syncify
+
+from ._group import group
+from ._logging import logger
+from ._task import P, R, TaskTemplate
+
+
+class MapAsyncIterator(Generic[P, R]):
+    """AsyncIterator implementation for the map function results"""
+
+    def __init__(self, func: TaskTemplate[P, R], args: tuple, name: str, concurrency: int, return_exceptions: bool):
+        self.func = func
+        self.args = args
+        self.name = name
+        self.concurrency = concurrency
+        self.return_exceptions = return_exceptions
+        self._tasks: List[asyncio.Task] = []
+        self._current_index = 0
+        self._completed_count = 0
+        self._exception_count = 0
+        self._task_count = 0
+        self._initialized = False
+
+    def __aiter__(self) -> AsyncIterator[Union[R, Exception]]:
+        """Return self as the async iterator"""
+        return self
+
+    async def __anext__(self) -> Union[R, Exception]:
+        """Get the next result"""
+        # Initialize on first call
+        if not self._initialized:
+            await self._initialize()
+
+        # Check if we've exhausted all tasks
+        if self._current_index >= self._task_count:
+            raise StopAsyncIteration
+
+        # Get the next task result
+        task = self._tasks[self._current_index]
+        self._current_index += 1
+
+        try:
+            result = await task
+            self._completed_count += 1
+            logger.debug(f"Task {self._current_index - 1} completed successfully")
+            return result
+        except Exception as e:
+            self._exception_count += 1
+            logger.debug(f"Task {self._current_index - 1} failed with exception: {e}")
+            if self.return_exceptions:
+                return e
+            else:
+                # Cancel remaining tasks
+                for remaining_task in self._tasks[self._current_index + 1 :]:
+                    remaining_task.cancel()
+                raise e
+
+    async def _initialize(self):
+        """Initialize the tasks - called lazily on first iteration"""
+        # Create all tasks at once
+        tasks = []
+        task_count = 0
+
+        for arg_tuple in zip(*self.args):
+            task = asyncio.create_task(self.func.aio(*arg_tuple))
+            tasks.append(task)
+            task_count += 1
+
+        if task_count == 0:
+            logger.info(f"Group '{self.name}' has no tasks to process")
+            self._tasks = []
+            self._task_count = 0
+        else:
+            logger.info(f"Starting {task_count} tasks in group '{self.name}' with unlimited concurrency")
+            self._tasks = tasks
+            self._task_count = task_count
+
+        self._initialized = True
+
+    async def collect(self) -> List[Union[R, Exception]]:
+        """Convenience method to collect all results into a list"""
+        results = []
+        async for result in self:
+            results.append(result)
+        return results
+
+    def __repr__(self):
+        return f"MapAsyncIterator(group_name='{self.name}', concurrency={self.concurrency})"
+
+
+class _Mapper(Generic[P, R]):
+    """
+    Internal mapper class to handle the mapping logic
+
+    NOTE: The reason why we do not use the `@syncify` decorator here is because, in `syncify` we cannot use
+    context managers like `group` directly in the function body. This is because the `__exit__` method of the
+    context manager is called after the function returns. An for `_context` the `__exit__` method releases the
+    token (for contextvar), which was created in a separate thread. This leads to an exception like:
+
+    """
+
+    @classmethod
+    def _get_name(cls, task_name: str, group_name: str | None) -> str:
+        """Get the name of the group, defaulting to 'map' if not provided."""
+        return f"{task_name}_{group_name or 'map'}"
+
+    def __call__(
+        self,
+        func: TaskTemplate[P, R],
+        *args: Iterable[Any],
+        group_name: str | None = None,
+        concurrency: int = 0,
+        return_exceptions: bool = True,
+    ) -> Iterator[Union[R, Exception]]:
+        """
+        Map a function over the provided arguments with concurrent execution.
+
+        :param func: The async function to map.
+        :param args: Positional arguments to pass to the function (iterables that will be zipped).
+        :param group_name: The name of the group for the mapped tasks.
+        :param concurrency: The maximum number of concurrent tasks to run. If 0, run all tasks concurrently.
+        :param return_exceptions: If True, yield exceptions instead of raising them.
+        :return: AsyncIterator yielding results in order.
+        """
+        if not args:
+            return
+
+        name = self._get_name(func.name, group_name)
+        logger.debug(f"Blocking Map for {name}")
+        with group(name):
+            import flyte
+
+            tctx = flyte.ctx()
+            if tctx is None or tctx.mode == "local":
+                logger.warning("Running map in local mode, which will run every task sequentially.")
+                for v in zip(*args):
+                    try:
+                        yield func(*v)  # type: ignore
+                    except Exception as e:
+                        if return_exceptions:
+                            yield e
+                        else:
+                            raise e
+                return
+
+            i = 0
+            for x in cast(
+                Iterator[R],
+                _map(
+                    func,
+                    *args,
+                    name=name,
+                    concurrency=concurrency,
+                    return_exceptions=True,
+                ),
+            ):
+                logger.debug(f"Mapped {x}, task {i}")
+                i += 1
+                yield x
+
+    async def aio(
+        self,
+        func: TaskTemplate[P, R],
+        *args: Iterable[Any],
+        group_name: str | None = None,
+        concurrency: int = 0,
+        return_exceptions: bool = True,
+    ) -> AsyncGenerator[Union[R, Exception], None]:
+        if not args:
+            return
+        name = self._get_name(func.name, group_name)
+        with group(name):
+            import flyte
+
+            tctx = flyte.ctx()
+            if tctx is None or tctx.mode == "local":
+                logger.warning("Running map in local mode, which will run every task sequentially.")
+                for v in zip(*args):
+                    try:
+                        yield func(*v)  # type: ignore
+                    except Exception as e:
+                        if return_exceptions:
+                            yield e
+                        else:
+                            raise e
+                return
+            async for x in _map.aio(
+                func,
+                *args,
+                name=name,
+                concurrency=concurrency,
+                return_exceptions=return_exceptions,
+            ):
+                yield cast(Union[R, Exception], x)
+
+
+@syncify
+async def _map(
+    func: TaskTemplate[P, R],
+    *args: Iterable[Any],
+    name: str = "map",
+    concurrency: int = 0,
+    return_exceptions: bool = True,
+) -> AsyncIterator[Union[R, Exception]]:
+    iter = MapAsyncIterator(
+        func=func, args=args, name=name, concurrency=concurrency, return_exceptions=return_exceptions
+    )
+    async for result in iter:
+        yield result
+
+
+map: _Mapper = _Mapper()

flyte/_run.py

@@ -1,5 +1,6 @@
 from __future__ import annotations
 
+import asyncio
 import pathlib
 import uuid
 from typing import TYPE_CHECKING, Any, Literal, Optional, Tuple, Union, cast
@@ -287,40 +288,39 @@ class _Runner:
 
     async def _run_local(self, obj: TaskTemplate[P, R], *args: P.args, **kwargs: P.kwargs) -> R:
         from flyte._internal.controllers import create_controller
-        from flyte._internal.runtime.convert import (
-            convert_error_to_native,
-            convert_from_native_to_inputs,
-            convert_outputs_to_native,
-        )
-        from flyte._internal.runtime.entrypoints import direct_dispatch
+        from flyte._internal.controllers._local_controller import LocalController
+        from flyte.report import Report
+
+        controller = cast(LocalController, create_controller("local"))
 
-        controller = create_controller("local")
-        inputs = await convert_from_native_to_inputs(obj.native_interface, *args, **kwargs)
         if self._name is None:
             action = ActionID.create_random()
         else:
             action = ActionID(name=self._name)
-        out, err = await direct_dispatch(
-            obj,
+
+        ctx = internal_ctx()
+        tctx = TaskContext(
             action=action,
-            raw_data_path=internal_ctx().raw_data,
-            version="na",
-            controller=controller,
-            inputs=inputs,
-            output_path=self._metadata_path,
-            run_base_dir=self._metadata_path,
             checkpoints=Checkpoints(
                 prev_checkpoint_path=internal_ctx().raw_data.path, checkpoint_path=internal_ctx().raw_data.path
             ),
-        )  # type: ignore
-        if err:
-            native_err = convert_error_to_native(err)
-            if native_err:
-                raise native_err
-        if obj.native_interface.outputs and len(obj.native_interface.outputs) > 0:
-            if out is not None:
-                return cast(R, await convert_outputs_to_native(obj.native_interface, out))
-        return cast(R, None)
+            code_bundle=None,
+            output_path=self._metadata_path,
+            run_base_dir=self._metadata_path,
+            version="na",
+            raw_data_path=internal_ctx().raw_data,
+            compiled_image_cache=None,
+            report=Report(name=action.name),
+            mode="local",
+        )
+        async with ctx.replace_task_context(tctx):
+            # make the local version always runs on a different thread, returns a wrapped future.
+            if obj._call_as_synchronous:
+                fut = controller.submit_sync(obj, *args, **kwargs)
+                awaitable = asyncio.wrap_future(fut)
+                return await awaitable
+            else:
+                return await controller.submit(obj, *args, **kwargs)
 
     @syncify
     async def run(self, task: TaskTemplate[P, Union[R, Run]], *args: P.args, **kwargs: P.kwargs) -> Union[R, Run]:
@@ -351,7 +351,7 @@ class _Runner:
             return await self._run_hybrid(task, *args, **kwargs)
 
         # TODO We could use this for remote as well and users could simply pass flyte:// or s3:// or file://
-        async with internal_ctx().new_raw_data_path(
+        with internal_ctx().new_raw_data_path(
             raw_data_path=RawDataPath.from_local_folder(local_folder=self._raw_data_path)
         ):
             return await self._run_local(task, *args, **kwargs)

flyte/_task.py

@@ -1,5 +1,6 @@
 from __future__ import annotations
 
+import asyncio
 import weakref
 from dataclasses import dataclass, field, replace
 from functools import cached_property
@@ -216,6 +217,51 @@ class TaskTemplate(Generic[P, R]):
     def native_interface(self) -> NativeInterface:
         return self.interface
 
+    async def aio(self, *args: P.args, **kwargs: P.kwargs) -> Coroutine[Any, Any, R] | R:
+        """
+        The aio function allows executing "sync" tasks, in an async context. This helps with migrating v1 defined sync
+        tasks to be used within an asyncio parent task.
+        This function will also re-raise exceptions from the underlying task.
+
+        Example:
+        ```python
+        @env.task
+        def my_legacy_task(x: int) -> int:
+            return x
+
+        @env.task
+        async def my_new_parent_task(n: int) -> List[int]:
+            collect = []
+            for x in range(n):
+                collect.append(my_legacy_task.aio(x))
+            return asyncio.gather(*collect)
+        ```
+        :param args:
+        :param kwargs:
+        :return:
+        """
+
+        ctx = internal_ctx()
+        if ctx.is_task_context():
+            from ._internal.controllers import get_controller
+
+            # If we are in a task context, that implies we are executing a Run.
+            # In this scenario, we should submit the task to the controller.
+            controller = get_controller()
+            if controller:
+                if self._call_as_synchronous:
+                    fut = controller.submit_sync(self, *args, **kwargs)
+                    asyncio_future = asyncio.wrap_future(fut)  # Wrap the future to make it awaitable
+                    return await asyncio_future
+                else:
+                    return await controller.submit(self, *args, **kwargs)
+            else:
+                raise RuntimeSystemError("BadContext", "Controller is not initialized.")
+        else:
+            # Local execute, just stay out of the way, but because .aio is used, we want to return an awaitable,
+            # even for synchronous tasks. This is to support migration.
+            return self.forward(*args, **kwargs)
+
     def __call__(self, *args: P.args, **kwargs: P.kwargs) -> Coroutine[Any, Any, R] | R:
         """
         This is the entrypoint for an async function task at runtime. It will be called during an execution.
@@ -235,15 +281,18 @@ class TaskTemplate(Generic[P, R]):
                 from ._internal.controllers import get_controller
 
                 controller = get_controller()
-                if controller:
-                    if self._call_as_synchronous:
-                        return controller.submit_sync(self, *args, **kwargs)
-                    else:
-                        return controller.submit(self, *args, **kwargs)
-                else:
+                if not controller:
                     raise RuntimeSystemError("BadContext", "Controller is not initialized.")
 
-            return self.forward(*args, **kwargs)
+                if self._call_as_synchronous:
+                    fut = controller.submit_sync(self, *args, **kwargs)
+                    x = fut.result(None)
+                    return x
+                else:
+                    return controller.submit(self, *args, **kwargs)
+            else:
+                # If not in task context, purely function run, stay out of the way
+                return self.forward(*args, **kwargs)
         except RuntimeSystemError:
             raise
         except RuntimeUserError:

flyte/_utils/helpers.py

@@ -1,6 +1,7 @@
 import os
 import string
 import typing
+from contextlib import contextmanager
 from pathlib import Path
 
 
@@ -106,3 +107,17 @@ def get_cwd_editable_install() -> typing.Optional[Path]:
                     return install  # note we want the install folder, not the parent
 
     return None
+
+
+@contextmanager
+def _selector_policy():
+    import asyncio
+
+    original_policy = asyncio.get_event_loop_policy()
+    try:
+        if os.name == "nt" and hasattr(asyncio, "WindowsSelectorEventLoopPolicy"):
+            asyncio.set_event_loop_policy(asyncio.WindowsSelectorEventLoopPolicy())
+
+        yield
+    finally:
+        asyncio.set_event_loop_policy(original_policy)

flyte/_version.py

@@ -17,5 +17,5 @@ __version__: str
 __version_tuple__: VERSION_TUPLE
 version_tuple: VERSION_TUPLE
 
-__version__ = version = '0.2.0b8'
-__version_tuple__ = version_tuple = (0, 2, 0, 'b8')
+__version__ = version = '0.2.0b9'
+__version_tuple__ = version_tuple = (0, 2, 0, 'b9')