flyte 2.0.0b18__py3-none-any.whl → 2.0.0b20__py3-none-any.whl

flyte/_bin/runtime.py

@@ -101,7 +101,6 @@ def main(
     from flyte._logging import logger
     from flyte.models import ActionID, Checkpoints, CodeBundle, RawDataPath
 
-    logger.warning(f"Flyte runtime started for action {name} with run name {run_name}")
     logger.info("Registering faulthandler for SIGUSR1")
     faulthandler.register(signal.SIGUSR1)
 
@@ -117,6 +116,8 @@ def main(
     if name.startswith("{{"):
         name = os.getenv("ACTION_NAME", "")
 
+    logger.warning(f"Flyte runtime started for action {name} with run name {run_name}")
+
     if debug and name == "a0":
         from flyte._debug.vscode import _start_vscode_server
 

flyte/_initialize.py

@@ -228,7 +228,7 @@ async def init(
 
 @syncify
 async def init_from_config(
-    path_or_config: str | Config | None = None,
+    path_or_config: str | Path | Config | None = None,
     root_dir: Path | None = None,
     log_level: int | None = None,
 ) -> None:
@@ -251,11 +251,11 @@ async def init_from_config(
     if path_or_config is None:
         # If no path is provided, use the default config file
         cfg = config.auto()
-    elif isinstance(path_or_config, str):
+    elif isinstance(path_or_config, (str, Path)):
         if root_dir:
-            cfg_path = str(root_dir / path_or_config)
+            cfg_path = root_dir.expanduser() / path_or_config
         else:
-            cfg_path = path_or_config
+            cfg_path = Path(path_or_config).expanduser()
         if not Path(cfg_path).exists():
             raise InitializationError(
                 "ConfigFileNotFoundError",

flyte/_internal/controllers/__init__.py

@@ -5,12 +5,13 @@ from typing import TYPE_CHECKING, Any, Callable, Literal, Optional, Protocol, Tu
 from flyte._task import TaskTemplate
 from flyte.models import ActionID, NativeInterface
 
+if TYPE_CHECKING:
+    from flyte.remote._task import TaskDetails
+
 from ._trace import TraceInfo
 
 __all__ = ["Controller", "ControllerType", "TraceInfo", "create_controller", "get_controller"]
 
-from ..._protos.workflow import task_definition_pb2
-
 if TYPE_CHECKING:
     import concurrent.futures
 
@@ -41,9 +42,7 @@ class Controller(Protocol):
         """
         ...
 
-    async def submit_task_ref(
-        self, _task: task_definition_pb2.TaskDetails, max_inline_io_bytes: int, *args, **kwargs
-    ) -> Any:
+    async def submit_task_ref(self, _task: "TaskDetails", *args, **kwargs) -> Any:
         """
         Submit a task reference to the controller asynchronously and wait for the result. This is async and will block
         the current coroutine until the result is available.

flyte/_internal/controllers/_local_controller.py

@@ -11,10 +11,10 @@ from flyte._internal.controllers import TraceInfo
 from flyte._internal.runtime import convert
 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
+from flyte.remote._task import TaskDetails
 
 R = TypeVar("R")
 
@@ -192,7 +192,7 @@ class LocalController:
         assert info.start_time
         assert info.end_time
 
-    async def submit_task_ref(
-        self, _task: task_definition_pb2.TaskDetails, max_inline_io_bytes: int, *args, **kwargs
-    ) -> Any:
-        raise flyte.errors.ReferenceTaskError("Reference tasks cannot be executed locally, only remotely.")
+    async def submit_task_ref(self, _task: TaskDetails, max_inline_io_bytes: int, *args, **kwargs) -> Any:
+        raise flyte.errors.ReferenceTaskError(
+            f"Reference tasks cannot be executed locally, only remotely. Found remote task {_task.name}"
+        )

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

@@ -12,7 +12,6 @@ from typing import Any, Awaitable, DefaultDict, Tuple, TypeVar
 import flyte
 import flyte.errors
 import flyte.storage as storage
-import flyte.types as types
 from flyte._code_bundle import build_pkl_bundle
 from flyte._context import internal_ctx
 from flyte._internal.controllers import TraceInfo
@@ -24,10 +23,11 @@ from flyte._internal.runtime.task_serde import translate_task_to_wire
 from flyte._internal.runtime.types_serde import transform_native_to_typed_interface
 from flyte._logging import logger
 from flyte._protos.common import identifier_pb2
-from flyte._protos.workflow import run_definition_pb2, task_definition_pb2
+from flyte._protos.workflow import run_definition_pb2
 from flyte._task import TaskTemplate
 from flyte._utils.helpers import _selector_policy
 from flyte.models import MAX_INLINE_IO_BYTES, ActionID, NativeInterface, SerializationContext
+from flyte.remote._task import TaskDetails
 
 R = TypeVar("R")
 
@@ -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)
@@ -482,19 +482,17 @@ class RemoteController(Controller):
                 # If the action is cancelled, we need to cancel the action on the server as well
                 raise
 
-    async def _submit_task_ref(
-        self, invoke_seq_num: int, _task: task_definition_pb2.TaskDetails, max_inline_io_bytes: int, *args, **kwargs
-    ) -> Any:
+    async def _submit_task_ref(self, invoke_seq_num: int, _task: TaskDetails, *args, **kwargs) -> Any:
         ctx = internal_ctx()
         tctx = ctx.data.task_context
         if tctx is None:
             raise flyte.errors.RuntimeSystemError("BadContext", "Task context not initialized")
         current_action_id = tctx.action
-        task_name = _task.spec.task_template.id.name
+        task_name = _task.name
+
+        native_interface = _task.interface
+        pb_interface = _task.pb2.spec.task_template.interface
 
-        native_interface = types.guess_interface(
-            _task.spec.task_template.interface, default_inputs=_task.spec.default_inputs
-        )
         inputs = await convert.convert_from_native_to_inputs(native_interface, *args, **kwargs)
         inputs_hash = convert.generate_inputs_hash_from_proto(inputs.proto_inputs)
         sub_action_id, sub_action_output_path = convert.generate_sub_action_id_and_output_path(
@@ -503,19 +501,19 @@ class RemoteController(Controller):
 
         serialized_inputs = inputs.proto_inputs.SerializeToString(deterministic=True)
         inputs_uri = io.inputs_path(sub_action_output_path)
-        await upload_inputs_with_retry(serialized_inputs, inputs_uri, max_inline_io_bytes)
+        await upload_inputs_with_retry(serialized_inputs, inputs_uri, _task.max_inline_io_bytes)
         # cache key - task name, task signature, inputs, cache version
         cache_key = None
-        md = _task.spec.task_template.metadata
+        md = _task.pb2.spec.task_template.metadata
         ignored_input_vars = []
         if len(md.cache_ignore_input_vars) > 0:
             ignored_input_vars = list(md.cache_ignore_input_vars)
-        if _task.spec.task_template.metadata and _task.spec.task_template.metadata.discoverable:
-            discovery_version = _task.spec.task_template.metadata.discovery_version
+        if md and md.discoverable:
+            discovery_version = md.discovery_version
             cache_key = convert.generate_cache_key_hash(
                 task_name,
                 inputs_hash,
-                _task.spec.task_template.interface,
+                pb_interface,
                 discovery_version,
                 ignored_input_vars,
                 inputs.proto_inputs,
@@ -537,7 +535,7 @@ class RemoteController(Controller):
             ),
             parent_action_name=current_action_id.name,
             group_data=tctx.group_data,
-            task_spec=_task.spec,
+            task_spec=_task.pb2.spec,
             inputs_uri=inputs_uri,
             run_output_base=tctx.run_base_dir,
             cache_key=cache_key,
@@ -566,12 +564,10 @@ class RemoteController(Controller):
                     "RuntimeError",
                     f"Task {n.action_id.name} did not return an output path, but the task has outputs defined.",
                 )
-            return await load_and_convert_outputs(native_interface, n.realized_outputs_uri, max_inline_io_bytes)
+            return await load_and_convert_outputs(native_interface, n.realized_outputs_uri, _task.max_inline_io_bytes)
         return None
 
-    async def submit_task_ref(
-        self, _task: task_definition_pb2.TaskDetails, max_inline_io_bytes: int, *args, **kwargs
-    ) -> Any:
+    async def submit_task_ref(self, _task: TaskDetails, *args, **kwargs) -> Any:
         ctx = internal_ctx()
         tctx = ctx.data.task_context
         if tctx is None:
@@ -579,4 +575,4 @@ class RemoteController(Controller):
         current_action_id = tctx.action
         task_call_seq = self.generate_task_call_sequence(_task, current_action_id)
         async with self._parent_action_semaphore[unique_action_name(current_action_id)]:
-            return await self._submit_task_ref(task_call_seq, _task, max_inline_io_bytes, *args, **kwargs)
+            return await self._submit_task_ref(task_call_seq, _task, *args, **kwargs)

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: