flyte 0.2.0b7__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/_environment.py

@@ -1,7 +1,8 @@
 from __future__ import annotations
 
+import re
 from dataclasses import dataclass, field
-from typing import TYPE_CHECKING, Dict, List, Literal, Optional, Union
+from typing import TYPE_CHECKING, Any, Dict, List, Literal, Optional, Union
 
 import rich.repr
 
@@ -14,6 +15,10 @@ if TYPE_CHECKING:
     from kubernetes.client import V1PodTemplate
 
 
+def is_snake_or_kebab_with_numbers(s: str) -> bool:
+    return re.fullmatch(r"^[a-z0-9]+([_-][a-z0-9]+)*$", s) is not None
+
+
 @rich.repr.auto
 @dataclass(init=True, repr=True)
 class Environment:
@@ -36,8 +41,44 @@ class Environment:
     resources: Optional[Resources] = None
     image: Union[str, Image, Literal["auto"]] = "auto"
 
+    def __post_init__(self):
+        if not is_snake_or_kebab_with_numbers(self.name):
+            raise ValueError(f"Environment name '{self.name}' must be in snake_case or kebab-case format.")
+
     def add_dependency(self, *env: Environment):
         """
         Add a dependency to the environment.
         """
         self.env_dep_hints.extend(env)
+
+    def clone_with(
+        self,
+        name: str,
+        image: Optional[Union[str, Image, Literal["auto"]]] = None,
+        resources: Optional[Resources] = None,
+        env: Optional[Dict[str, str]] = None,
+        secrets: Optional[SecretRequest] = None,
+        env_dep_hints: Optional[List[Environment]] = None,
+        **kwargs: Any,
+    ) -> Environment:
+        raise NotImplementedError
+
+    def _get_kwargs(self) -> Dict[str, Any]:
+        """
+        Get the keyword arguments for the environment.
+        """
+        kwargs: Dict[str, Any] = {
+            "env_dep_hints": self.env_dep_hints,
+            "image": self.image,
+        }
+        if self.resources is not None:
+            kwargs["resources"] = self.resources
+        if self.secrets is not None:
+            kwargs["secrets"] = self.secrets
+        if self.env is not None:
+            kwargs["env"] = self.env
+        if self.pod_template is not None:
+            kwargs["pod_template"] = self.pod_template
+        if self.description is not None:
+            kwargs["description"] = self.description
+        return kwargs

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/_image.py

@@ -444,8 +444,7 @@ class Image:
         ```
 
         For more information on the uv script format, see the documentation:
-        <href="https://docs.astral.sh/uv/guides/scripts/#declaring-script-dependencies">
-            UV: Declaring script dependencies</href>
+        [UV: Declaring script dependencies](https://docs.astral.sh/uv/guides/scripts/#declaring-script-dependencies)
 
         :param name: name of the image
         :param registry: registry to use for the image

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,6 +32,15 @@ class Controller(Protocol):
         """
         ...
 
+    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:
         """
         Submit a task reference to the controller asynchronously and wait for the result. This is async and will block

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,19 +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.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
@@ -58,6 +111,18 @@ class LocalController:
             return result
         return out
 
+    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,6 +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.helpers import _selector_policy
 from flyte.models import ActionID, NativeInterface, SerializationContext
 
 R = TypeVar("R")
@@ -119,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:
         """
@@ -235,6 +241,48 @@ 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)
 
+    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):
         """
         This method is invoked when the parent action is finished. It will finalize the run and upload the outputs

flyte/_internal/controllers/remote/_informer.py

@@ -38,11 +38,11 @@ class ActionCache:
         """
         Add an action to the cache if it doesn't exist. This is invoked by the watch.
         """
-        logger.info(f"Observing phase {run_definition_pb2.Phase.Name(state.phase)} for {state.action_id.name}")
+        logger.debug(f"Observing phase {run_definition_pb2.Phase.Name(state.phase)} for {state.action_id.name}")
         if state.output_uri:
-            logger.info(f"Output URI: {state.output_uri}")
+            logger.debug(f"Output URI: {state.output_uri}")
         else:
-            logger.info(f"{state.action_id.name} has no output URI")
+            logger.warning(f"{state.action_id.name} has no output URI")
         if state.phase == run_definition_pb2.Phase.PHASE_FAILED:
             logger.error(
                 f"Action {state.action_id.name} failed with error (msg):"

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)