flyte 2.0.0b19__py3-none-any.whl → 2.0.0b21__py3-none-any.whl

flyte/_bin/runtime.py

@@ -21,8 +21,8 @@ import click
 
 ACTION_NAME = "ACTION_NAME"
 RUN_NAME = "RUN_NAME"
-PROJECT_NAME = "FLYTE_INTERNAL_TASK_PROJECT"
-DOMAIN_NAME = "FLYTE_INTERNAL_TASK_DOMAIN"
+PROJECT_NAME = "FLYTE_INTERNAL_EXECUTION_PROJECT"
+DOMAIN_NAME = "FLYTE_INTERNAL_EXECUTION_DOMAIN"
 ORG_NAME = "_U_ORG_NAME"
 ENDPOINT_OVERRIDE = "_U_EP_OVERRIDE"
 RUN_OUTPUT_BASE_DIR = "_U_RUN_BASE"

flyte/_code_bundle/_ignore.py

@@ -83,8 +83,15 @@ class StandardIgnore(Ignore):
         self.patterns = patterns if patterns else STANDARD_IGNORE_PATTERNS
 
     def _is_ignored(self, path: pathlib.Path) -> bool:
+        # Convert to relative path for pattern matching
+        try:
+            rel_path = path.relative_to(self.root)
+        except ValueError:
+            # If path is not under root, don't ignore it
+            return False
+
         for pattern in self.patterns:
-            if fnmatch(str(path), pattern):
+            if fnmatch(str(rel_path), pattern):
                 return True
         return False
 
@@ -105,9 +112,10 @@ class IgnoreGroup(Ignore):
 
     def list_ignored(self) -> List[str]:
         ignored = []
-        for dir, _, files in self.root.walk():
+        for dir, _, files in os.walk(self.root):
+            dir_path = Path(dir)
             for file in files:
-                abs_path = dir / file
+                abs_path = dir_path / file
                 if self.is_ignored(abs_path):
                     ignored.append(str(abs_path.relative_to(self.root)))
         return ignored

flyte/_code_bundle/_packaging.py

@@ -14,10 +14,9 @@ import typing
 from typing import List, Optional, Tuple, Union
 
 import click
-from rich import print as rich_print
 from rich.tree import Tree
 
-from flyte._logging import logger
+from flyte._logging import _get_console, logger
 
 from ._ignore import Ignore, IgnoreGroup
 from ._utils import CopyFiles, _filehash_update, _pathhash_update, ls_files, tar_strip_file_attributes
@@ -27,10 +26,10 @@ FAST_FILEENDING = ".tar.gz"
 
 
 def print_ls_tree(source: os.PathLike, ls: typing.List[str]):
-    click.secho("Files to be copied for fast registration...", fg="bright_blue")
+    logger.info("Files to be copied for fast registration...")
 
     tree_root = Tree(
-        f":open_file_folder: [link file://{source}]{source} (detected source root)",
+        f"File structure:\n:open_file_folder: {source}",
         guide_style="bold bright_blue",
     )
     trees = {pathlib.Path(source): tree_root}
@@ -49,7 +48,12 @@ def print_ls_tree(source: os.PathLike, ls: typing.List[str]):
                 else:
                     current = trees[current_path]
         trees[fpp.parent].add(f"{fpp.name}", guide_style="bold bright_blue")
-    rich_print(tree_root)
+
+    console = _get_console()
+    with console.capture() as capture:
+        console.print(tree_root, overflow="ignore", no_wrap=True, crop=False)
+    logger.info(f"Root directory: [link=file://{source}]{source}[/link]")
+    logger.info(capture.get(), extra={"console": console})
 
 
 def _compress_tarball(source: pathlib.Path, output: pathlib.Path) -> None:

flyte/_code_bundle/_utils.py

@@ -156,7 +156,7 @@ def list_all_files(source_path: pathlib.Path, deref_symlinks, ignore_group: Opti
 
     # This is needed to prevent infinite recursion when walking with followlinks
     visited_inodes = set()
-    for root, dirnames, files in source_path.walk(top_down=True, follow_symlinks=deref_symlinks):
+    for root, dirnames, files in os.walk(source_path, topdown=True, followlinks=deref_symlinks):
         dirnames[:] = [d for d in dirnames if d not in EXCLUDE_DIRS]
         if deref_symlinks:
             inode = os.stat(root).st_ino
@@ -167,7 +167,7 @@ def list_all_files(source_path: pathlib.Path, deref_symlinks, ignore_group: Opti
         ff = []
         files.sort()
         for fname in files:
-            abspath = (root / fname).absolute()
+            abspath = (pathlib.Path(root) / fname).absolute()
             # Only consider files that exist (e.g. disregard symlinks that point to non-existent files)
             if not os.path.exists(abspath):
                 logger.info(f"Skipping non-existent file {abspath}")

flyte/_initialize.py

@@ -110,6 +110,12 @@ async def _initialize_client(
     )
 
 
+def _initialize_logger(log_level: int | None = None):
+    initialize_logger(enable_rich=True)
+    if log_level:
+        initialize_logger(log_level=log_level, enable_rich=True)
+
+
 @syncify
 async def init(
     org: str | None = None,
@@ -172,14 +178,9 @@ async def init(
 
     :return: None
     """
-    from flyte._tools import ipython_check
     from flyte._utils import get_cwd_editable_install, org_from_endpoint, sanitize_endpoint
 
-    interactive_mode = ipython_check()
-
-    initialize_logger(enable_rich=interactive_mode)
-    if log_level:
-        initialize_logger(log_level=log_level, enable_rich=interactive_mode)
+    _initialize_logger(log_level=log_level)
 
     global _init_config  # noqa: PLW0603
 
@@ -231,6 +232,7 @@ async def init_from_config(
     path_or_config: str | Path | Config | None = None,
     root_dir: Path | None = None,
     log_level: int | None = None,
+    storage: Storage | None = None,
 ) -> None:
     """
     Initialize the Flyte system using a configuration file or Config object. This method should be called before any
@@ -245,6 +247,8 @@ async def init_from_config(
         default is set using the default initialization policies
     :return: None
     """
+    from rich.highlighter import ReprHighlighter
+
     import flyte.config as config
 
     cfg: config.Config
@@ -266,7 +270,9 @@ async def init_from_config(
     else:
         cfg = path_or_config
 
-    logger.debug(f"Flyte config initialized as {cfg}")
+    _initialize_logger(log_level=log_level)
+
+    logger.info(f"Flyte config initialized as {cfg}", extra={"highlighter": ReprHighlighter()})
     await init.aio(
         org=cfg.task.org,
         project=cfg.task.project,
@@ -283,6 +289,7 @@ async def init_from_config(
         root_dir=root_dir,
         log_level=log_level,
         image_builder=cfg.image.builder,
+        storage=storage,
     )
 
 

flyte/_interface.py

@@ -1,7 +1,11 @@
 from __future__ import annotations
 
 import inspect
-from typing import Dict, Generator, Tuple, Type, TypeVar, Union, cast, get_args, get_type_hints
+import typing
+from enum import Enum
+from typing import Dict, Generator, Literal, Tuple, Type, TypeVar, Union, cast, get_args, get_origin, get_type_hints
+
+from flyte._logging import logger
 
 
 def default_output_name(index: int = 0) -> str:
@@ -69,7 +73,15 @@ def extract_return_annotation(return_annotation: Union[Type, Tuple, None]) -> Di
         if len(return_annotation.__args__) == 1:  # type: ignore
             raise TypeError("Tuples should be used to indicate multiple return values, found only one return variable.")
         ra = get_args(return_annotation)
-        return dict(zip(list(output_name_generator(len(ra))), ra))
+        annotations = {}
+        for i, r in enumerate(ra):
+            if r is Ellipsis:
+                raise TypeError("Variable length tuples are not supported as return types.")
+            if get_origin(r) is Literal:
+                annotations[default_output_name(i)] = literal_to_enum(cast(Type, r))
+            else:
+                annotations[default_output_name(i)] = r
+        return annotations
 
     elif isinstance(return_annotation, tuple):
         if len(return_annotation) == 1:
@@ -79,4 +91,25 @@ def extract_return_annotation(return_annotation: Union[Type, Tuple, None]) -> Di
     else:
         # Handle all other single return types
         # Task returns unnamed native tuple
+        if get_origin(return_annotation) is Literal:
+            return {default_output_name(): literal_to_enum(cast(Type, return_annotation))}
         return {default_output_name(): cast(Type, return_annotation)}
+
+
+def literal_to_enum(literal_type: Type) -> Type[Enum | typing.Any]:
+    """Convert a Literal[...] into Union[str, Enum]."""
+
+    if get_origin(literal_type) is not Literal:
+        raise TypeError(f"{literal_type} is not a Literal")
+
+    values = get_args(literal_type)
+    if not all(isinstance(v, str) for v in values):
+        logger.warning(f"Literal type {literal_type} contains non-string values, using Any instead of Enum")
+        return typing.Any
+    # Deduplicate & keep order
+    enum_dict = {str(v).upper(): v for v in values}
+
+    # Dynamically create an Enum
+    literal_enum = Enum("LiteralEnum", enum_dict)  # type: ignore
+
+    return literal_enum  # type: ignore

flyte/_internal/controllers/remote/__init__.py

@@ -54,7 +54,5 @@ def create_remote_controller(
 
     controller = RemoteController(
         client_coro=client_coro,
-        workers=10,
-        max_system_retries=5,
     )
     return controller

flyte/_internal/controllers/remote/_controller.py

@@ -117,9 +117,8 @@ class RemoteController(Controller):
     def __init__(
         self,
         client_coro: Awaitable[ClientSet],
-        workers: int,
-        max_system_retries: int,
-        default_parent_concurrency: int = 100,
+        workers: int = 20,
+        max_system_retries: int = 10,
     ):
         """ """
         super().__init__(
@@ -127,6 +126,7 @@ class RemoteController(Controller):
             workers=workers,
             max_system_retries=max_system_retries,
         )
+        default_parent_concurrency = int(os.getenv("_F_P_CNC", "100"))
         self._default_parent_concurrency = default_parent_concurrency
         self._parent_action_semaphore: DefaultDict[str, asyncio.Semaphore] = defaultdict(
             lambda: asyncio.Semaphore(default_parent_concurrency)

flyte/_internal/controllers/remote/_core.py

@@ -1,12 +1,14 @@
 from __future__ import annotations
 
 import asyncio
+import os
 import sys
 import threading
 from asyncio import Event
 from typing import Awaitable, Coroutine, Optional
 
 import grpc.aio
+from aiolimiter import AsyncLimiter
 from google.protobuf.wrappers_pb2 import StringValue
 
 import flyte.errors
@@ -32,10 +34,10 @@ class Controller:
     def __init__(
         self,
         client_coro: Awaitable[ClientSet],
-        workers: int = 2,
-        max_system_retries: int = 5,
+        workers: int = 20,
+        max_system_retries: int = 10,
         resource_log_interval_sec: float = 10.0,
-        min_backoff_on_err_sec: float = 0.1,
+        min_backoff_on_err_sec: float = 0.5,
         thread_wait_timeout_sec: float = 5.0,
         enqueue_timeout_sec: float = 5.0,
     ):
@@ -53,14 +55,17 @@ class Controller:
         self._running = False
         self._resource_log_task = None
         self._workers = workers
-        self._max_retries = max_system_retries
+        self._max_retries = int(os.getenv("_F_MAX_RETRIES", max_system_retries))
         self._resource_log_interval = resource_log_interval_sec
         self._min_backoff_on_err = min_backoff_on_err_sec
+        self._max_backoff_on_err = float(os.getenv("_F_MAX_BFF_ON_ERR", "10.0"))
         self._thread_wait_timeout = thread_wait_timeout_sec
         self._client_coro = client_coro
         self._failure_event: Event | None = None
         self._enqueue_timeout = enqueue_timeout_sec
         self._informer_start_wait_timeout = thread_wait_timeout_sec
+        max_qps = int(os.getenv("_F_MAX_QPS", "100"))
+        self._rate_limiter = AsyncLimiter(max_qps, 1.0)
 
         # Thread management
         self._thread = None
@@ -194,15 +199,16 @@ class Controller:
         # We will wait for this to signal that the thread is ready
         # Signal the main thread that we're ready
         logger.debug("Background thread initialization complete")
-        self._thread_ready.set()
         if sys.version_info >= (3, 11):
             async with asyncio.TaskGroup() as tg:
                 for i in range(self._workers):
-                    tg.create_task(self._bg_run())
+                    tg.create_task(self._bg_run(f"worker-{i}"))
+                self._thread_ready.set()
         else:
             tasks = []
             for i in range(self._workers):
-                tasks.append(asyncio.create_task(self._bg_run()))
+                tasks.append(asyncio.create_task(self._bg_run(f"worker-{i}")))
+            self._thread_ready.set()
             await asyncio.gather(*tasks)
 
     def _bg_thread_target(self):
@@ -221,6 +227,7 @@ class Controller:
         except Exception as e:
             logger.error(f"Controller thread encountered an exception: {e}")
             self._set_exception(e)
+            self._failure_event.set()
         finally:
             if self._loop and self._loop.is_running():
                 self._loop.close()
@@ -292,21 +299,22 @@ class Controller:
         started = action.is_started()
         action.mark_cancelled()
         if started:
-            logger.info(f"Cancelling action: {action.name}")
-            try:
-                # TODO add support when the queue service supports aborting actions
-                # await self._queue_service.AbortQueuedAction(
-                #     queue_service_pb2.AbortQueuedActionRequest(action_id=action.action_id),
-                #     wait_for_ready=True,
-                # )
-                logger.info(f"Successfully cancelled action: {action.name}")
-            except grpc.aio.AioRpcError as e:
-                if e.code() in [
-                    grpc.StatusCode.NOT_FOUND,
-                    grpc.StatusCode.FAILED_PRECONDITION,
-                ]:
-                    logger.info(f"Action {action.name} not found, assumed completed or cancelled.")
-                    return
+            async with self._rate_limiter:
+                logger.info(f"Cancelling action: {action.name}")
+                try:
+                    # TODO add support when the queue service supports aborting actions
+                    # await self._queue_service.AbortQueuedAction(
+                    #     queue_service_pb2.AbortQueuedActionRequest(action_id=action.action_id),
+                    #     wait_for_ready=True,
+                    # )
+                    logger.info(f"Successfully cancelled action: {action.name}")
+                except grpc.aio.AioRpcError as e:
+                    if e.code() in [
+                        grpc.StatusCode.NOT_FOUND,
+                        grpc.StatusCode.FAILED_PRECONDITION,
+                    ]:
+                        logger.info(f"Action {action.name} not found, assumed completed or cancelled.")
+                        return
         else:
             # If the action is not started, we have to ensure it does not get launched
             logger.info(f"Action {action.name} is not started, no need to cancel.")
@@ -320,56 +328,69 @@ class Controller:
         Attempt to launch an action.
         """
         if not action.is_started():
-            task: queue_service_pb2.TaskAction | None = None
-            trace: queue_service_pb2.TraceAction | None = None
-            if action.type == "task":
-                if action.task is None:
-                    raise flyte.errors.RuntimeSystemError(
-                        "NoTaskSpec", "Task Spec not found, cannot launch Task Action."
+            async with self._rate_limiter:
+                task: queue_service_pb2.TaskAction | None = None
+                trace: queue_service_pb2.TraceAction | None = None
+                if action.type == "task":
+                    if action.task is None:
+                        raise flyte.errors.RuntimeSystemError(
+                            "NoTaskSpec", "Task Spec not found, cannot launch Task Action."
+                        )
+                    cache_key = None
+                    logger.info(f"Action {action.name} has cache version {action.cache_key}")
+                    if action.cache_key:
+                        cache_key = StringValue(value=action.cache_key)
+
+                    task = queue_service_pb2.TaskAction(
+                        id=task_definition_pb2.TaskIdentifier(
+                            version=action.task.task_template.id.version,
+                            org=action.task.task_template.id.org,
+                            project=action.task.task_template.id.project,
+                            domain=action.task.task_template.id.domain,
+                            name=action.task.task_template.id.name,
+                        ),
+                        spec=action.task,
+                        cache_key=cache_key,
                     )
-                cache_key = None
-                logger.info(f"Action {action.name} has cache version {action.cache_key}")
-                if action.cache_key:
-                    cache_key = StringValue(value=action.cache_key)
-
-                task = queue_service_pb2.TaskAction(
-                    id=task_definition_pb2.TaskIdentifier(
-                        version=action.task.task_template.id.version,
-                        org=action.task.task_template.id.org,
-                        project=action.task.task_template.id.project,
-                        domain=action.task.task_template.id.domain,
-                        name=action.task.task_template.id.name,
-                    ),
-                    spec=action.task,
-                    cache_key=cache_key,
-                )
-            elif action.type == "trace":
-                trace = action.trace
-
-            logger.debug(f"Attempting to launch action: {action.name}")
-            try:
-                await self._queue_service.EnqueueAction(
-                    queue_service_pb2.EnqueueActionRequest(
-                        action_id=action.action_id,
-                        parent_action_name=action.parent_action_name,
-                        task=task,
-                        trace=trace,
-                        input_uri=action.inputs_uri,
-                        run_output_base=action.run_output_base,
-                        group=action.group.name if action.group else None,
-                        # Subject is not used in the current implementation
-                    ),
-                    wait_for_ready=True,
-                    timeout=self._enqueue_timeout,
-                )
-                logger.info(f"Successfully launched action: {action.name}")
-            except grpc.aio.AioRpcError as e:
-                if e.code() == grpc.StatusCode.ALREADY_EXISTS:
-                    logger.info(f"Action {action.name} already exists, continuing to monitor.")
-                    return
-                logger.exception(f"Failed to launch action: {action.name} backing off...")
-                logger.debug(f"Action details: {action}")
-                raise e
+                elif action.type == "trace":
+                    trace = action.trace
+
+                logger.debug(f"Attempting to launch action: {action.name}")
+                try:
+                    await self._queue_service.EnqueueAction(
+                        queue_service_pb2.EnqueueActionRequest(
+                            action_id=action.action_id,
+                            parent_action_name=action.parent_action_name,
+                            task=task,
+                            trace=trace,
+                            input_uri=action.inputs_uri,
+                            run_output_base=action.run_output_base,
+                            group=action.group.name if action.group else None,
+                            # Subject is not used in the current implementation
+                        ),
+                        wait_for_ready=True,
+                        timeout=self._enqueue_timeout,
+                    )
+                    logger.info(f"Successfully launched action: {action.name}")
+                except grpc.aio.AioRpcError as e:
+                    if e.code() == grpc.StatusCode.ALREADY_EXISTS:
+                        logger.info(f"Action {action.name} already exists, continuing to monitor.")
+                        return
+                    if e.code() in [
+                        grpc.StatusCode.FAILED_PRECONDITION,
+                        grpc.StatusCode.INVALID_ARGUMENT,
+                        grpc.StatusCode.NOT_FOUND,
+                    ]:
+                        raise flyte.errors.RuntimeSystemError(
+                            e.code().name, f"Precondition failed: {e.details()}"
+                        ) from e
+                    # For all other errors, we will retry with backoff
+                    logger.exception(
+                        f"Failed to launch action: {action.name}, Code: {e.code()}, "
+                        f"Details {e.details()} backing off..."
+                    )
+                    logger.debug(f"Action details: {action}")
+                    raise flyte.errors.SlowDownError(f"Failed to launch action: {e.details()}") from e
 
     @log
     async def _bg_process(self, action: Action):
@@ -397,35 +418,42 @@ class Controller:
             await asyncio.sleep(self._resource_log_interval)
 
     @log
-    async def _bg_run(self):
+    async def _bg_run(self, worker_id: str):
         """Run loop with resource status logging"""
+        logger.info(f"Worker {worker_id} started")
         while self._running:
             logger.debug(f"{threading.current_thread().name} Waiting for resource")
             action = await self._shared_queue.get()
             logger.debug(f"{threading.current_thread().name} Got resource {action.name}")
             try:
                 await self._bg_process(action)
-            except Exception as e:
-                logger.error(f"Error in controller loop: {e}")
-                # TODO we need a better way of handling backoffs currently the entire worker coroutine backs off
-                await asyncio.sleep(self._min_backoff_on_err)
-                action.increment_retries()
+            except flyte.errors.SlowDownError as e:
+                action.retries += 1
                 if action.retries > self._max_retries:
-                    err = flyte.errors.RuntimeSystemError(
-                        code=type(e).__name__,
-                        message=f"Controller failed, system retries {action.retries}"
-                        f" crossed threshold {self._max_retries}",
-                    )
-                    err.__cause__ = e
-                    action.set_client_error(err)
-                    informer = await self._informers.get(
-                        run_name=action.run_name,
-                        parent_action_name=action.parent_action_name,
-                    )
-                    if informer:
-                        await informer.fire_completion_event(action.name)
-                else:
-                    await self._shared_queue.put(action)
+                    raise
+                backoff = min(self._min_backoff_on_err * (2 ** (action.retries - 1)), self._max_backoff_on_err)
+                logger.warning(
+                    f"[{worker_id}] Backing off for {backoff} [retry {action.retries}/{self._max_retries}] "
+                    f"on action {action.name} due to error: {e}"
+                )
+                await asyncio.sleep(backoff)
+                logger.warning(f"[{worker_id}] Retrying action {action.name} after backoff")
+                await self._shared_queue.put(action)
+            except Exception as e:
+                logger.error(f"[{worker_id}] Error in controller loop: {e}")
+                err = flyte.errors.RuntimeSystemError(
+                    code=type(e).__name__,
+                    message=f"Controller failed, system retries {action.retries} crossed threshold {self._max_retries}",
+                    worker=worker_id,
+                )
+                err.__cause__ = e
+                action.set_client_error(err)
+                informer = await self._informers.get(
+                    run_name=action.run_name,
+                    parent_action_name=action.parent_action_name,
+                )
+                if informer:
+                    await informer.fire_completion_event(action.name)
             finally:
                 self._shared_queue.task_done()
 

flyte/_internal/controllers/remote/_informer.py

@@ -132,8 +132,10 @@ class Informer:
         parent_action_name: str,
         shared_queue: Queue,
         client: Optional[StateService] = None,
-        watch_backoff_interval_sec: float = 1.0,
+        min_watch_backoff: float = 1.0,
+        max_watch_backoff: float = 30.0,
         watch_conn_timeout_sec: float = 5.0,
+        max_watch_retries: int = 10,
     ):
         self.name = self.mkname(run_name=run_id.name, parent_action_name=parent_action_name)
         self.parent_action_name = parent_action_name
@@ -144,8 +146,10 @@ class Informer:
         self._running = False
         self._watch_task: asyncio.Task | None = None
         self._ready = asyncio.Event()
-        self._watch_backoff_interval_sec = watch_backoff_interval_sec
+        self._min_watch_backoff = min_watch_backoff
+        self._max_watch_backoff = max_watch_backoff
         self._watch_conn_timeout_sec = watch_conn_timeout_sec
+        self._max_watch_retries = max_watch_retries
 
     @classmethod
     def mkname(cls, *, run_name: str, parent_action_name: str) -> str:
@@ -211,13 +215,16 @@ class Informer:
         """
         # sentinel = False
         retries = 0
-        max_retries = 5
         last_exc = None
         while self._running:
-            if retries >= max_retries:
-                logger.error(f"Informer watch failure retries crossed threshold {retries}/{max_retries}, exiting!")
+            if retries >= self._max_watch_retries:
+                logger.error(
+                    f"Informer watch failure retries crossed threshold {retries}/{self._max_watch_retries}, exiting!"
+                )
                 raise last_exc
             try:
+                if retries >= 1:
+                    logger.warning(f"Informer watch retrying, attempt {retries}/{self._max_watch_retries}")
                 watcher = self._client.Watch(
                     state_service_pb2.WatchRequest(
                         parent_action_id=identifier_pb2.ActionIdentifier(
@@ -252,7 +259,9 @@ class Informer:
                 logger.exception(f"Watch error: {self.name}", exc_info=e)
                 last_exc = e
                 retries += 1
-            await asyncio.sleep(self._watch_backoff_interval_sec)
+            backoff = min(self._min_watch_backoff * (2**retries), self._max_watch_backoff)
+            logger.warning(f"Watch for {self.name} failed, retrying in {backoff} seconds...")
+            await asyncio.sleep(backoff)
 
     @log
     async def start(self, timeout: Optional[float] = None) -> asyncio.Task: