flyte 0.2.0b1__py3-none-any.whl → 2.0.0b46__py3-none-any.whl

flyte/_internal/controllers/remote/_core.py

@@ -1,17 +1,21 @@
 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 flyteidl2.common import identifier_pb2
+from flyteidl2.task import task_definition_pb2
+from flyteidl2.workflow import queue_service_pb2, run_definition_pb2
+from google.protobuf.wrappers_pb2 import StringValue
 
 import flyte.errors
 from flyte._logging import log, logger
-from flyte._protos.workflow import queue_service_pb2, run_definition_pb2, task_definition_pb2
-from flyte.errors import RuntimeSystemError
 
 from ._action import Action
 from ._informer import InformerCache
@@ -28,11 +32,11 @@ 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,
-        thread_wait_timeout_sec: float = 10.0,
+        min_backoff_on_err_sec: float = 0.5,
+        thread_wait_timeout_sec: float = 5.0,
         enqueue_timeout_sec: float = 5.0,
     ):
         """
@@ -49,14 +53,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
@@ -79,21 +86,19 @@ class Controller:
         """Public API to submit a resource and wait for completion"""
         return await self._run_coroutine_in_controller_thread(self._bg_submit_action(action))
 
-    async def get_action(
-        self, action_id: run_definition_pb2.ActionIdentifier, parent_action_name: str
-    ) -> Optional[Action]:
+    async def get_action(self, action_id: identifier_pb2.ActionIdentifier, parent_action_name: str) -> Optional[Action]:
         """Get the action from the informer"""
-        informer = await self._informers.get(run_name=action_id.run.name, parent_action_name=parent_action_name)
-        if informer:
-            return await informer.get(action_id.name)
-        return None
+        return await self._run_coroutine_in_controller_thread(self._bg_get_action(action_id, parent_action_name))
 
     @log
     async def cancel_action(self, action: Action):
         return await self._run_coroutine_in_controller_thread(self._bg_cancel_action(action))
 
     async def _finalize_parent_action(
-        self, run_id: run_definition_pb2.RunIdentifier, parent_action_name: str, timeout: Optional[float] = None
+        self,
+        run_id: identifier_pb2.RunIdentifier,
+        parent_action_name: str,
+        timeout: Optional[float] = None,
     ):
         """Finalize the parent run"""
         await self._run_coroutine_in_controller_thread(
@@ -111,19 +116,21 @@ class Controller:
                     raise RuntimeError("Failure event not initialized")
                 self._failure_event.set()
         except asyncio.CancelledError:
-            pass
+            raise
 
     async def _bg_watch_for_errors(self):
         if self._failure_event is None:
             raise RuntimeError("Failure event not initialized")
         await self._failure_event.wait()
         logger.warning(f"Failure event received: {self._failure_event}, cleaning up informers and exiting.")
+        self._running = False
 
     async def watch_for_errors(self):
         """Watch for errors in the background thread"""
         await self._run_coroutine_in_controller_thread(self._bg_watch_for_errors())
-        raise RuntimeSystemError(
-            code="InformerWatchFailure", message=f"Controller thread failed with exception: {self._get_exception()}"
+        raise flyte.errors.RuntimeSystemError(
+            code="InformerWatchFailure",
+            message=f"Controller thread failed with exception: {self._get_exception()}",
         )
 
     @log
@@ -142,7 +149,6 @@ class Controller:
         with self._thread_com_lock:
             return self._thread_exception
 
-    @log
     def _start(self):
         """Start the controller in a separate thread"""
         if self._thread and self._thread.is_alive():
@@ -155,13 +161,14 @@ class Controller:
         self._thread.start()
 
         # Wait for the thread to be ready
-        logger.info("Waiting for controller thread to be ready...")
         if not self._thread_ready.wait(timeout=self._thread_wait_timeout):
+            logger.warning("Controller thread did not finish within timeout")
             raise TimeoutError("Controller thread failed to start in time")
 
         if self._get_exception():
-            raise RuntimeSystemError(
-                type(self._get_exception()).__name__, f"Controller thread startup failed: {self._get_exception()}"
+            raise flyte.errors.RuntimeSystemError(
+                type(self._get_exception()).__name__,
+                f"Controller thread startup failed: {self._get_exception()}",
             )
 
         logger.info(f"Controller started in thread: {self._thread.name}")
@@ -190,15 +197,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):
@@ -207,6 +215,7 @@ class Controller:
             # Create a new event loop for this thread
             self._loop = asyncio.new_event_loop()
             asyncio.set_event_loop(self._loop)
+            self._loop.set_exception_handler(flyte.errors.silence_grpc_polling_error)
             logger.debug(f"Controller thread started with new event loop: {threading.current_thread().name}")
 
             # Create an event to signal the errors were observed in the thread's loop
@@ -216,13 +225,14 @@ 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()
             logger.debug(f"Controller thread exiting: {threading.current_thread().name}")
 
     async def _bg_get_action(
-        self, action_id: run_definition_pb2.ActionIdentifier, parent_action_name: str
+        self, action_id: identifier_pb2.ActionIdentifier, parent_action_name: str
     ) -> Optional[Action]:
         """Get the action from the informer"""
         # Ensure the informer is created and wait for it to be ready
@@ -239,13 +249,15 @@ class Controller:
         return None
 
     async def _bg_finalize_informer(
-        self, run_id: run_definition_pb2.RunIdentifier, parent_action_name: str, timeout: Optional[float] = None
+        self,
+        run_id: identifier_pb2.RunIdentifier,
+        parent_action_name: str,
+        timeout: Optional[float] = None,
     ):
         informer = await self._informers.remove(run_name=run_id.name, parent_action_name=parent_action_name)
         if informer:
             await informer.stop()
 
-    @log
     async def _bg_submit_action(self, action: Action) -> Action:
         """Submit a resource and await its completion, returning the final state"""
         logger.debug(f"{threading.current_thread().name} Submitting action {action.name}")
@@ -270,7 +282,7 @@ class Controller:
             raise ValueError(f"Action {action.name} not found")
         logger.debug(f"{threading.current_thread().name} Removed completion event for action {action.name}")
         await informer.remove(action.name)  # TODO we should not remove maybe, we should keep a record of completed?
-        logger.debug(f"{threading.current_thread().name} Removed action {action.name}, final={final_resource}")
+        logger.debug(f"{threading.current_thread().name} Removed action {action.name}")
         return final_resource
 
     async def _bg_cancel_action(self, action: Action):
@@ -284,17 +296,21 @@ class Controller:
         started = action.is_started()
         action.mark_cancelled()
         if started:
-            logger.info(f"Cancelling action: {action.name}")
-            try:
-                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:
+                    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.")
@@ -307,41 +323,72 @@ class Controller:
         """
         Attempt to launch an action.
         """
-        if not action.is_started() and action.task is not None:
-            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=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,
+        if not action.is_started():
+            async with self._rate_limiter:
+                task: run_definition_pb2.TaskAction | None = None
+                trace: run_definition_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 = run_definition_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,
                         ),
-                        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
+                        spec=action.task,
+                        cache_key=cache_key,
+                        cluster=action.queue,
+                    )
+                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):
         """Process resource updates"""
         logger.debug(f"Processing action: name={action.name}, started={action.is_started()}")
@@ -358,39 +405,52 @@ class Controller:
     async def _bg_log_stats(self):
         """Periodically log resource stats if debug is enabled"""
         while self._running:
-            async for started, pending, terminal in self._informers.count_started_pending_terminal_actions():
+            async for (
+                started,
+                pending,
+                terminal,
+            ) in self._informers.count_started_pending_terminal_actions():
                 logger.info(f"Resource stats: Started={started}, Pending={pending}, Terminal={terminal}")
             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()
-                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
+                try:
+                    await self._bg_process(action)
+                except flyte.errors.SlowDownError as e:
+                    action.retries += 1
+                    if action.retries > self._max_retries:
+                        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}"
                     )
-                    if informer:
-                        await informer.fire_completion_event(action.name)
-                else:
+                    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 for {action.name}: {e}")
+                err = flyte.errors.RuntimeSystemError(
+                    code=type(e).__name__,
+                    message=f"Controller failed, system retries {action.retries} / {self._max_retries} "
+                    f"crossed threshold, for action {action.name}: {e}",
+                    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

@@ -5,9 +5,10 @@ from asyncio import Queue
 from typing import AsyncIterator, Callable, Dict, Optional, Tuple, cast
 
 import grpc.aio
+from flyteidl2.common import identifier_pb2, phase_pb2
+from flyteidl2.workflow import state_service_pb2
 
 from flyte._logging import log, logger
-from flyte._protos.workflow import run_definition_pb2, state_service_pb2
 
 from ._action import Action
 from ._service_protocol import StateService
@@ -38,12 +39,14 @@ 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 {phase_pb2.ActionPhase.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")
-        if state.phase == run_definition_pb2.Phase.PHASE_FAILED:
+            logger.warning(
+                f"{state.action_id.name} has no output URI, in phase {phase_pb2.ActionPhase.Name(state.phase)}"
+            )
+        if state.phase == phase_pb2.ACTION_PHASE_FAILED:
             logger.error(
                 f"Action {state.action_id.name} failed with error (msg):"
                 f" [{state.error if state.HasField('error') else None}]"
@@ -125,12 +128,14 @@ class Informer:
 
     def __init__(
         self,
-        run_id: run_definition_pb2.RunIdentifier,
+        run_id: identifier_pb2.RunIdentifier,
         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
@@ -141,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:
@@ -208,16 +215,19 @@ 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=run_definition_pb2.ActionIdentifier(
+                        parent_action_id=identifier_pb2.ActionIdentifier(
                             name=self.parent_action_name,
                             run=self._run_id,
                         ),
@@ -235,7 +245,7 @@ class Informer:
                     await self._shared_queue.put(node)
                     # hack to work in the absence of sentinel
             except asyncio.CancelledError:
-                logger.warning(f"Watch cancelled: {self.name}")
+                logger.info(f"Watch cancelled: {self.name}")
                 return
             except asyncio.TimeoutError as e:
                 logger.error(f"Watch timeout: {self.name}", exc_info=e)
@@ -249,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:
@@ -258,7 +270,7 @@ class Informer:
             logger.warning("Informer already running")
             return cast(asyncio.Task, self._watch_task)
         self._running = True
-        self._watch_task = asyncio.create_task(self.watch())
+        self._watch_task = asyncio.create_task(self.watch(), name=f"InformerWatch-{self.parent_action_name}")
         await self.wait_for_cache_sync(timeout=timeout)
         return self._watch_task
 
@@ -288,7 +300,7 @@ class InformerCache:
     @log
     async def get_or_create(
         self,
-        run_id: run_definition_pb2.RunIdentifier,
+        run_id: identifier_pb2.RunIdentifier,
         parent_action_name: str,
         shared_queue: Queue,
         state_service: StateService,
@@ -330,20 +342,28 @@ class InformerCache:
     async def get(self, *, run_name: str, parent_action_name: str) -> Informer | None:
         """Get an informer by name"""
         async with self._lock:
-            return self._cache.get(Informer.mkname(run_name=run_name, parent_action_name=parent_action_name), None)
+            return self._cache.get(
+                Informer.mkname(run_name=run_name, parent_action_name=parent_action_name),
+                None,
+            )
 
     @log
     async def remove(self, *, run_name: str, parent_action_name: str) -> Informer | None:
         """Remove an informer from the cache"""
         async with self._lock:
-            return self._cache.pop(Informer.mkname(run_name=run_name, parent_action_name=parent_action_name), None)
+            return self._cache.pop(
+                Informer.mkname(run_name=run_name, parent_action_name=parent_action_name),
+                None,
+            )
 
     async def has(self, *, run_name: str, parent_action_name: str) -> bool:
         """Check if an informer exists in the cache"""
         async with self._lock:
             return Informer.mkname(run_name=run_name, parent_action_name=parent_action_name) in self._cache
 
-    async def count_started_pending_terminal_actions(self) -> AsyncIterator[Tuple[int, int, int]]:
+    async def count_started_pending_terminal_actions(
+        self,
+    ) -> AsyncIterator[Tuple[int, int, int]]:
         """Log resource stats"""
         async with self._lock:
             for informer in self._cache.values():
@@ -353,7 +373,7 @@ class InformerCache:
         """Stop all informers and remove them from the cache"""
         async with self._lock:
             while self._cache:
-                name, informer = self._cache.popitem()
+                _name, informer = self._cache.popitem()
                 try:
                     await informer.stop()
                 except asyncio.CancelledError:

flyte/_internal/controllers/remote/_service_protocol.py

@@ -2,7 +2,7 @@ from __future__ import annotations
 
 from typing import AsyncIterator, Protocol
 
-from flyte._protos.workflow import queue_service_pb2, state_service_pb2
+from flyteidl2.workflow import queue_service_pb2, state_service_pb2
 
 
 class StateService(Protocol):
@@ -33,7 +33,7 @@ class QueueService(Protocol):
         req: queue_service_pb2.AbortQueuedActionRequest,
         **kwargs,
     ) -> queue_service_pb2.AbortQueuedActionResponse:
-        """Dequeue a task"""
+        """Cancel an enqueued task"""
 
 
 class ClientSet(Protocol):

flyte/_internal/imagebuild/__init__.py

@@ -1,11 +1,3 @@
-import asyncio
-from typing import List
+from flyte._internal.imagebuild.image_builder import ImageBuildEngine
 
-from flyte._image import Image
-from flyte._internal.imagebuild.docker_builder import DockerImageBuilder
-
-
-async def build(images: List[Image]) -> List[str]:
-    builder = DockerImageBuilder()
-    ts = [asyncio.create_task(builder.build_image(image)) for image in images]
-    return list(await asyncio.gather(*ts))
+__all__ = ["ImageBuildEngine"]