flyte 2.0.0b22__py3-none-any.whl → 2.0.0b30__py3-none-any.whl

flyte/_interface.py

@@ -51,6 +51,8 @@ def extract_return_annotation(return_annotation: Union[Type, Tuple, None]) -> Di
     Note that Options 1 and 2 are identical, just syntactic sugar. In the NamedTuple case, we'll use the names in the
     definition. In all other cases, we'll automatically generate output names, indexed starting at 0.
     """
+    if isinstance(return_annotation, str):
+        raise TypeError("String return annotations are not supported.")
 
     # Handle Option 6
     # We can think about whether we should add a default output name with type None in the future.

flyte/_internal/controllers/_local_controller.py

@@ -2,16 +2,20 @@ import asyncio
 import atexit
 import concurrent.futures
 import os
+import pathlib
 import threading
 from typing import Any, Callable, Tuple, TypeVar
 
 import flyte.errors
+from flyte._cache.cache import VersionParameters, cache_from_request
+from flyte._cache.local_cache import LocalTaskCache
 from flyte._context import internal_ctx
 from flyte._internal.controllers import TraceInfo
 from flyte._internal.runtime import convert
 from flyte._internal.runtime.entrypoints import direct_dispatch
+from flyte._internal.runtime.types_serde import transform_native_to_typed_interface
 from flyte._logging import log, logger
-from flyte._task import TaskTemplate
+from flyte._task import AsyncFunctionTaskTemplate, TaskTemplate
 from flyte._utils.helpers import _selector_policy
 from flyte.models import ActionID, NativeInterface
 from flyte.remote._task import TaskDetails
@@ -81,31 +85,67 @@ class LocalController:
             raise flyte.errors.RuntimeSystemError("BadContext", "Task context not initialized")
 
         inputs = await convert.convert_from_native_to_inputs(_task.native_interface, *args, **kwargs)
-        serialized_inputs = inputs.proto_inputs.SerializeToString(deterministic=True)
+        inputs_hash = convert.generate_inputs_hash_from_proto(inputs.proto_inputs)
+        task_interface = transform_native_to_typed_interface(_task.interface)
 
         sub_action_id, sub_action_output_path = convert.generate_sub_action_id_and_output_path(
-            tctx, _task.name, serialized_inputs, 0
+            tctx, _task.name, inputs_hash, 0
         )
         sub_action_raw_data_path = tctx.raw_data_path
-
-        out, err = await direct_dispatch(
-            _task,
-            controller=self,
-            action=sub_action_id,
-            raw_data_path=sub_action_raw_data_path,
-            inputs=inputs,
-            version=tctx.version,
-            checkpoints=tctx.checkpoints,
-            code_bundle=tctx.code_bundle,
-            output_path=sub_action_output_path,
-            run_base_dir=tctx.run_base_dir,
+        # Make sure the output path exists
+        pathlib.Path(sub_action_output_path).mkdir(parents=True, exist_ok=True)
+        pathlib.Path(sub_action_raw_data_path.path).mkdir(parents=True, exist_ok=True)
+
+        task_cache = cache_from_request(_task.cache)
+        cache_enabled = task_cache.is_enabled()
+        if isinstance(_task, AsyncFunctionTaskTemplate):
+            version_parameters = VersionParameters(func=_task.func, image=_task.image)
+        else:
+            version_parameters = VersionParameters(func=None, image=_task.image)
+        cache_version = task_cache.get_version(version_parameters)
+        cache_key = convert.generate_cache_key_hash(
+            _task.name,
+            inputs_hash,
+            task_interface,
+            cache_version,
+            list(task_cache.get_ignored_inputs()),
+            inputs.proto_inputs,
         )
-        if err:
-            exc = convert.convert_error_to_native(err)
-            if exc:
-                raise exc
-            else:
-                raise flyte.errors.RuntimeSystemError("BadError", "Unknown error")
+
+        out = None
+        # We only get output from cache if the cache behavior is set to auto
+        if task_cache.behavior == "auto":
+            out = await LocalTaskCache.get(cache_key)
+            if out is not None:
+                logger.info(
+                    f"Cache hit for task '{_task.name}' (version: {cache_version}), getting result from cache..."
+                )
+
+        if out is None:
+            out, err = await direct_dispatch(
+                _task,
+                controller=self,
+                action=sub_action_id,
+                raw_data_path=sub_action_raw_data_path,
+                inputs=inputs,
+                version=cache_version,
+                checkpoints=tctx.checkpoints,
+                code_bundle=tctx.code_bundle,
+                output_path=sub_action_output_path,
+                run_base_dir=tctx.run_base_dir,
+            )
+
+            if err:
+                exc = convert.convert_error_to_native(err)
+                if exc:
+                    raise exc
+                else:
+                    raise flyte.errors.RuntimeSystemError("BadError", "Unknown error")
+
+            # store into cache
+            if cache_enabled and out is not None:
+                await LocalTaskCache.set(cache_key, out)
+
         if _task.native_interface.outputs:
             if out is None:
                 raise flyte.errors.RuntimeSystemError("BadOutput", "Task output not captured.")
@@ -129,7 +169,7 @@ class LocalController:
         pass
 
     async def stop(self):
-        pass
+        await LocalTaskCache.close()
 
     async def watch_for_errors(self):
         pass
@@ -146,16 +186,17 @@ class LocalController:
         tctx = ctx.data.task_context
         if not tctx:
             raise flyte.errors.NotInTaskContextError("BadContext", "Task context not initialized")
+
         converted_inputs = convert.Inputs.empty()
         if _interface.inputs:
             converted_inputs = await convert.convert_from_native_to_inputs(_interface, *args, **kwargs)
             assert converted_inputs
 
-        serialized_inputs = converted_inputs.proto_inputs.SerializeToString(deterministic=True)
+        inputs_hash = convert.generate_inputs_hash_from_proto(converted_inputs.proto_inputs)
         action_id, action_output_path = convert.generate_sub_action_id_and_output_path(
             tctx,
             _func.__name__,
-            serialized_inputs,
+            inputs_hash,
             0,
         )
         assert action_output_path

flyte/_internal/controllers/_trace.py

@@ -1,7 +1,7 @@
 from dataclasses import dataclass, field
 from typing import Any, Optional
 
-from flyteidl.core import interface_pb2
+from flyteidl2.core import interface_pb2
 
 from flyte.models import ActionID, NativeInterface
 

flyte/_internal/controllers/remote/_action.py

@@ -1,18 +1,17 @@
 from __future__ import annotations
 
 from dataclasses import dataclass
-from typing import Literal
+from typing import Literal, Optional
 
-from flyteidl.core import execution_pb2, interface_pb2
-from google.protobuf import timestamp_pb2
-
-from flyte._protos.common import identifier_pb2
-from flyte._protos.workflow import (
-    queue_service_pb2,
+from flyteidl2.common import identifier_pb2
+from flyteidl2.core import execution_pb2, interface_pb2
+from flyteidl2.task import common_pb2, task_definition_pb2
+from flyteidl2.workflow import (
     run_definition_pb2,
     state_service_pb2,
-    task_definition_pb2,
 )
+from google.protobuf import timestamp_pb2
+
 from flyte.models import GroupData
 
 ActionType = Literal["task", "trace"]
@@ -31,7 +30,7 @@ class Action:
     friendly_name: str | None = None
     group: GroupData | None = None
     task: task_definition_pb2.TaskSpec | None = None
-    trace: queue_service_pb2.TraceAction | None = None
+    trace: run_definition_pb2.TraceAction | None = None
     inputs_uri: str | None = None
     run_output_base: str | None = None
     realized_outputs_uri: str | None = None
@@ -39,6 +38,7 @@ class Action:
     phase: run_definition_pb2.Phase | None = None
     started: bool = False
     retries: int = 0
+    queue: Optional[str] = None  # The queue to which this action was submitted.
     client_err: Exception | None = None  # This error is set when something goes wrong in the controller.
     cache_key: str | None = None  # None means no caching, otherwise it is the version of the cache.
 
@@ -122,6 +122,7 @@ class Action:
         inputs_uri: str,
         run_output_base: str,
         cache_key: str | None = None,
+        queue: Optional[str] = None,
     ) -> Action:
         return cls(
             action_id=sub_action_id,
@@ -132,6 +133,7 @@ class Action:
             inputs_uri=inputs_uri,
             run_output_base=run_output_base,
             cache_key=cache_key,
+            queue=queue,
         )
 
     @classmethod
@@ -195,12 +197,12 @@ class Action:
             realized_outputs_uri=outputs_uri,
             phase=run_definition_pb2.Phase.PHASE_SUCCEEDED,
             run_output_base=run_output_base,
-            trace=queue_service_pb2.TraceAction(
+            trace=run_definition_pb2.TraceAction(
                 name=friendly_name,
                 phase=run_definition_pb2.Phase.PHASE_SUCCEEDED,
                 start_time=st,
                 end_time=et,
-                outputs=run_definition_pb2.OutputReferences(
+                outputs=common_pb2.OutputReferences(
                     output_uri=outputs_uri,
                     report_uri=report_uri,
                 ),

flyte/_internal/controllers/remote/_client.py

@@ -1,8 +1,8 @@
 from __future__ import annotations
 
 import grpc.aio
+from flyteidl2.workflow import queue_service_pb2_grpc, state_service_pb2_grpc
 
-from flyte._protos.workflow import queue_service_pb2_grpc, state_service_pb2_grpc
 from flyte.remote import create_channel
 
 from ._service_protocol import QueueService, StateService

flyte/_internal/controllers/remote/_controller.py

@@ -9,6 +9,9 @@ from collections.abc import Callable
 from pathlib import Path
 from typing import Any, Awaitable, DefaultDict, Tuple, TypeVar
 
+from flyteidl2.common import identifier_pb2
+from flyteidl2.workflow import run_definition_pb2
+
 import flyte
 import flyte.errors
 import flyte.storage as storage
@@ -22,8 +25,6 @@ from flyte._internal.runtime import convert, io
 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
 from flyte._task import TaskTemplate
 from flyte._utils.helpers import _selector_policy
 from flyte.models import MAX_INLINE_IO_BYTES, ActionID, NativeInterface, SerializationContext
@@ -126,7 +127,7 @@ class RemoteController(Controller):
             workers=workers,
             max_system_retries=max_system_retries,
         )
-        default_parent_concurrency = int(os.getenv("_F_P_CNC", "100"))
+        default_parent_concurrency = int(os.getenv("_F_P_CNC", "1000"))
         self._default_parent_concurrency = default_parent_concurrency
         self._parent_action_semaphore: DefaultDict[str, asyncio.Semaphore] = defaultdict(
             lambda: asyncio.Semaphore(default_parent_concurrency)
@@ -238,6 +239,7 @@ class RemoteController(Controller):
             inputs_uri=inputs_uri,
             run_output_base=tctx.run_base_dir,
             cache_key=cache_key,
+            queue=_task.queue,
         )
 
         try:
@@ -375,11 +377,13 @@ class RemoteController(Controller):
 
         func_name = _func.__name__
         invoke_seq_num = self.generate_task_call_sequence(_func, current_action_id)
+
         inputs = await convert.convert_from_native_to_inputs(_interface, *args, **kwargs)
         serialized_inputs = inputs.proto_inputs.SerializeToString(deterministic=True)
+        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(
-            tctx, func_name, serialized_inputs, invoke_seq_num
+            tctx, func_name, inputs_hash, invoke_seq_num
         )
 
         inputs_uri = io.inputs_path(sub_action_output_path)
@@ -539,6 +543,7 @@ class RemoteController(Controller):
             inputs_uri=inputs_uri,
             run_output_base=tctx.run_base_dir,
             cache_key=cache_key,
+            queue=None,
         )
 
         try:

flyte/_internal/controllers/remote/_core.py

@@ -9,15 +9,13 @@ 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.common import identifier_pb2
-from flyte._protos.workflow import (
-    queue_service_pb2,
-    task_definition_pb2,
-)
 
 from ._action import Action
 from ._informer import InformerCache
@@ -118,13 +116,14 @@ 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"""
@@ -302,11 +301,10 @@ class Controller:
             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,
-                    # )
+                    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 [
@@ -329,8 +327,8 @@ class Controller:
         """
         if not action.is_started():
             async with self._rate_limiter:
-                task: queue_service_pb2.TaskAction | None = None
-                trace: queue_service_pb2.TraceAction | None = None
+                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(
@@ -341,7 +339,7 @@ class Controller:
                     if action.cache_key:
                         cache_key = StringValue(value=action.cache_key)
 
-                    task = queue_service_pb2.TaskAction(
+                    task = run_definition_pb2.TaskAction(
                         id=task_definition_pb2.TaskIdentifier(
                             version=action.task.task_template.id.version,
                             org=action.task.task_template.id.org,
@@ -351,6 +349,7 @@ class Controller:
                         ),
                         spec=action.task,
                         cache_key=cache_key,
+                        cluster=action.queue,
                     )
                 elif action.type == "trace":
                     trace = action.trace
@@ -440,10 +439,11 @@ class Controller:
                 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}")
+                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} crossed threshold {self._max_retries}",
+                    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

flyte/_internal/controllers/remote/_informer.py

@@ -5,10 +5,10 @@ from asyncio import Queue
 from typing import AsyncIterator, Callable, Dict, Optional, Tuple, cast
 
 import grpc.aio
+from flyteidl2.common import identifier_pb2
+from flyteidl2.workflow import run_definition_pb2, state_service_pb2
 
 from flyte._logging import log, logger
-from flyte._protos.common import identifier_pb2
-from flyte._protos.workflow import run_definition_pb2, state_service_pb2
 
 from ._action import Action
 from ._service_protocol import StateService
@@ -270,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
 
@@ -373,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):
@@ -28,12 +28,12 @@ class QueueService(Protocol):
     ) -> queue_service_pb2.EnqueueActionResponse:
         """Enqueue a task"""
 
-    # async def AbortQueuedAction(
-    #     self,
-    #     req: queue_service_pb2.AbortQueuedActionRequest,
-    #     **kwargs,
-    # ) -> queue_service_pb2.AbortQueuedActionResponse:
-    #     """Dequeue a task"""
+    async def AbortQueuedAction(
+        self,
+        req: queue_service_pb2.AbortQueuedActionRequest,
+        **kwargs,
+    ) -> queue_service_pb2.AbortQueuedActionResponse:
+        """Cancel an enqueued task"""
 
 
 class ClientSet(Protocol):