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

flyte/__init__.py

@@ -38,7 +38,8 @@ __all__ = [
     "deploy",
     "group",
     "init",
-    "init_auto_from_config",
+    "init_from_config",
+    "map",
     "run",
     "trace",
     "with_runcontext",
@@ -50,7 +51,8 @@ from ._deploy import deploy
 from ._environment import Environment
 from ._group import group
 from ._image import Image
-from ._initialize import init, init_auto_from_config
+from ._initialize import init, init_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/_deploy.py

@@ -128,6 +128,9 @@ async def apply(deployment: DeploymentPlan, copy_style: CopyFiles, dryrun: bool
     else:
         code_bundle = await build_code_bundle(from_dir=cfg.root_dir, dryrun=dryrun, copy_style=copy_style)
         deployment.version = code_bundle.computed_version
+        # TODO we should update the version to include the image cache digest and code bundle digest. This is
+        # to ensure that changes in image dependencies, cause an update to the deployment version.
+        # TODO Also hash the environment and tasks to ensure that changes in the environment or tasks
 
     sc = SerializationContext(
         project=cfg.project,

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

@@ -138,14 +138,13 @@ async def init(
     :param project: Optional project name (not used in this implementation)
     :param domain: Optional domain name (not used in this implementation)
     :param root_dir: Optional root directory from which to determine how to load files, and find paths to files.
+      This is useful for determining the root directory for the current project, and for locating files like config etc.
+      also use to determine all the code that needs to be copied to the remote location.
       defaults to the editable install directory if the cwd is in a Python editable install, else just the cwd.
     :param log_level: Optional logging level for the logger, default is set using the default initialization policies
     :param api_key: Optional API key for authentication
     :param endpoint: Optional API endpoint URL
     :param headless: Optional Whether to run in headless mode
-    :param mode: Optional execution model (local, remote). Default is local. When local is used,
-           the execution will be done locally. When remote is used, the execution will be sent to a remote server,
-           In the remote case, the endpoint or api_key must be set.
     :param insecure_skip_verify: Whether to skip SSL certificate verification
     :param auth_client_config: Optional client configuration for authentication
     :param auth_type: The authentication type to use (Pkce, ClientSecret, ExternalCommand, DeviceFlow)
@@ -177,6 +176,9 @@ async def init(
 
     global _init_config  # noqa: PLW0603
 
+    if endpoint and "://" not in endpoint:
+        endpoint = f"dns:///{endpoint}"
+
     with _init_lock:
         client = None
         if endpoint or api_key:
@@ -209,12 +211,16 @@ async def init(
 
 
 @syncify
-async def init_auto_from_config(path_or_config: str | Config | None = None) -> None:
+async def init_from_config(path_or_config: str | Config | None = None, root_dir: Path | None = None) -> None:
     """
     Initialize the Flyte system using a configuration file or Config object. This method should be called before any
     other Flyte remote API methods are called. Thread-safe implementation.
 
     :param path_or_config: Path to the configuration file or Config object
+    :param root_dir: Optional root directory from which to determine how to load files, and find paths to
+        files like config etc. For example if one uses the copy-style=="all", it is essential to determine the
+        root directory for the current project. If not provided, it defaults to the editable install directory or
+        if not available, the current working directory.
     :return: None
     """
     import flyte.config as config
@@ -222,7 +228,10 @@ async def init_auto_from_config(path_or_config: str | Config | None = None) -> N
     cfg: config.Config
     if path_or_config is None or isinstance(path_or_config, str):
         # If a string is passed, treat it as a path to the config file
-        cfg = config.auto(path_or_config)
+        if root_dir and path_or_config:
+            cfg = config.auto(str(root_dir / path_or_config))
+        else:
+            cfg = config.auto(path_or_config)
     else:
         # If a Config object is passed, use it directly
         cfg = path_or_config
@@ -241,6 +250,7 @@ async def init_auto_from_config(path_or_config: str | Config | None = None) -> N
         proxy_command=cfg.platform.proxy_command,
         client_id=cfg.platform.client_id,
         client_credentials_secret=cfg.platform.client_credentials_secret,
+        root_dir=root_dir,
     )
 
 

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.models import ActionID, NativeInterface, RawDataPath
+from flyte._utils.helpers import _selector_policy
+from flyte.models import ActionID, NativeInterface
 
 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
@@ -34,7 +86,7 @@ class LocalController:
         sub_action_id, sub_action_output_path = convert.generate_sub_action_id_and_output_path(
             tctx, _task.name, serialized_inputs, 0
         )
-        sub_action_raw_data_path = RawDataPath(path=sub_action_output_path)
+        sub_action_raw_data_path = tctx.raw_data_path
 
         out, err = await direct_dispatch(
             _task,
@@ -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()