flyte 0.0.1b3__py3-none-any.whl → 0.2.0a0__py3-none-any.whl

flyte/_internal/controllers/_local_controller.py

@@ -1,26 +1,79 @@
-from typing import Any, Tuple, TypeVar
+import asyncio
+import atexit
+import concurrent.futures
+import os
+import threading
+from typing import Any, Callable, Tuple, TypeVar
 
 import flyte.errors
 from flyte._context import internal_ctx
-from flyte._datastructures import ActionID, NativeInterface, RawDataPath
 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
 
 R = TypeVar("R")
 
 
+class _TaskRunner:
+    """A task runner that runs an asyncio event loop on a background thread."""
+
+    def __init__(self) -> None:
+        self.__loop: asyncio.AbstractEventLoop | None = None
+        self.__runner_thread: threading.Thread | None = None
+        self.__lock = threading.Lock()
+        atexit.register(self._close)
+
+    def _close(self) -> None:
+        if self.__loop:
+            self.__loop.stop()
+
+    def _execute(self) -> None:
+        loop = self.__loop
+        assert loop is not None
+        try:
+            loop.run_forever()
+        finally:
+            loop.close()
+
+    def get_exc_handler(self):
+        def exc_handler(loop, context):
+            logger.error(
+                f"Taskrunner for {self.__runner_thread.name if self.__runner_thread else 'no thread'} caught"
+                f" exception in {loop}: {context}"
+            )
+
+        return exc_handler
+
+    def get_run_future(self, coro: Any) -> concurrent.futures.Future:
+        """Synchronously run a coroutine on a background thread."""
+        name = f"{threading.current_thread().name} : loop-runner"
+        with self.__lock:
+            if self.__loop is None:
+                with _selector_policy():
+                    self.__loop = asyncio.new_event_loop()
+
+                exc_handler = self.get_exc_handler()
+                self.__loop.set_exception_handler(exc_handler)
+                self.__runner_thread = threading.Thread(target=self._execute, daemon=True, name=name)
+                self.__runner_thread.start()
+        fut = asyncio.run_coroutine_threadsafe(coro, self.__loop)
+        return fut
+
+
 class LocalController:
     def __init__(self):
         logger.debug("LocalController init")
+        self._runner_map: dict[str, _TaskRunner] = {}
 
     @log
     async def submit(self, _task: TaskTemplate, *args, **kwargs) -> Any:
         """
-        Submit a node to the controller
+        Main entrypoint for submitting a task to the local controller.
         """
         ctx = internal_ctx()
         tctx = ctx.data.task_context
@@ -28,8 +81,12 @@ class LocalController:
             raise flyte.errors.RuntimeSystemError("BadContext", "Task context not initialized")
 
         inputs = await convert.convert_from_native_to_inputs(_task.native_interface, *args, **kwargs)
-        sub_action_id, sub_action_output_path = convert.generate_sub_action_id_and_output_path(tctx, _task.name, inputs)
-        sub_action_raw_data_path = RawDataPath(path=sub_action_output_path)
+        serialized_inputs = inputs.proto_inputs.SerializeToString(deterministic=True)
+
+        sub_action_id, sub_action_output_path = convert.generate_sub_action_id_and_output_path(
+            tctx, _task.name, serialized_inputs, 0
+        )
+        sub_action_raw_data_path = tctx.raw_data_path
 
         out, err = await direct_dispatch(
             _task,
@@ -54,6 +111,18 @@ class LocalController:
             return result
         return out
 
+    def submit_sync(self, _task: TaskTemplate, *args, **kwargs) -> concurrent.futures.Future:
+        name = threading.current_thread().name + f"PID:{os.getpid()}"
+        coro = self.submit(_task, *args, **kwargs)
+        if name not in self._runner_map:
+            if len(self._runner_map) > 100:
+                logger.warning(
+                    "More than 100 event loop runners created!!! This could be a case of runaway recursion..."
+                )
+            self._runner_map[name] = _TaskRunner()
+
+        return self._runner_map[name].get_run_future(coro)
+
     async def finalize_parent_action(self, action: ActionID):
         pass
 
@@ -64,7 +133,7 @@ class LocalController:
         pass
 
     async def get_action_outputs(
-        self, _interface: NativeInterface, _func_name: str, *args, **kwargs
+        self, _interface: NativeInterface, _func: Callable, *args, **kwargs
     ) -> Tuple[TraceInfo, bool]:
         """
         This method returns the outputs of the action, if it is available.
@@ -79,8 +148,13 @@ class LocalController:
         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)
         action_id, action_output_path = convert.generate_sub_action_id_and_output_path(
-            tctx, _func_name, converted_inputs
+            tctx,
+            _func.__name__,
+            serialized_inputs,
+            0,
         )
         assert action_output_path
         return (
@@ -88,6 +162,7 @@ class LocalController:
                 action=action_id,
                 interface=_interface,
                 inputs_path=action_output_path,
+                name=_func.__name__,
             ),
             True,
         )
@@ -105,7 +180,7 @@ class LocalController:
 
         if info.interface.outputs and info.output:
             # If the result is not an AsyncGenerator, convert it directly
-            converted_outputs = await convert.convert_from_native_to_outputs(info.output, info.interface)
+            converted_outputs = await convert.convert_from_native_to_outputs(info.output, info.interface, info.name)
             assert converted_outputs
         elif info.error:
             # If there is an error, convert it to a native error

flyte/_internal/controllers/_trace.py

@@ -2,7 +2,7 @@ from dataclasses import dataclass
 from datetime import timedelta
 from typing import Any, Optional
 
-from flyte._datastructures import ActionID, NativeInterface
+from flyte.models import ActionID, NativeInterface
 
 
 @dataclass
@@ -18,6 +18,7 @@ class TraceInfo:
     duration: Optional[timedelta] = None
     output: Optional[Any] = None
     error: Optional[Exception] = None
+    name: str = ""
 
     def add_outputs(self, output: Any, duration: timedelta):
         """

flyte/_internal/controllers/remote/__init__.py

@@ -10,13 +10,13 @@ __all__ = ["RemoteController", "create_remote_controller"]
 def create_remote_controller(
     *,
     api_key: str | None = None,
-    auth_type: AuthType = "Pkce",
-    endpoint: str,
-    client_config: ClientConfig | None = None,
-    headless: bool = False,
+    endpoint: str | None = None,
     insecure: bool = False,
     insecure_skip_verify: bool = False,
     ca_cert_file_path: str | None = None,
+    client_config: ClientConfig | None = None,
+    auth_type: AuthType = "Pkce",
+    headless: bool = False,
     command: List[str] | None = None,
     proxy_command: List[str] | None = None,
     client_id: str | None = None,
@@ -27,13 +27,33 @@ def create_remote_controller(
     """
     Create a new instance of the remote controller.
     """
+    assert endpoint or api_key, "Either endpoint or api_key must be provided when initializing remote controller"
     from ._client import ControllerClient
     from ._controller import RemoteController
 
+    if endpoint:
+        client_coro = ControllerClient.for_endpoint(
+            endpoint,
+            insecure=insecure,
+            insecure_skip_verify=insecure_skip_verify,
+            ca_cert_file_path=ca_cert_file_path,
+            client_id=client_id,
+            client_credentials_secret=client_credentials_secret,
+            auth_type=auth_type,
+        )
+    elif api_key:
+        client_coro = ControllerClient.for_api_key(
+            api_key,
+            insecure=insecure,
+            insecure_skip_verify=insecure_skip_verify,
+            ca_cert_file_path=ca_cert_file_path,
+            client_id=client_id,
+            client_credentials_secret=client_credentials_secret,
+            auth_type=auth_type,
+        )
+
     controller = RemoteController(
-        client_coro=ControllerClient.for_endpoint(
-            endpoint=endpoint, insecure=insecure, insecure_skip_verify=insecure_skip_verify
-        ),
+        client_coro=client_coro,
         workers=10,
         max_system_retries=5,
     )

flyte/_internal/controllers/remote/_action.py

@@ -4,8 +4,8 @@ from dataclasses import dataclass
 
 from flyteidl.core import execution_pb2
 
-from flyte._datastructures import GroupData
 from flyte._protos.workflow import run_definition_pb2, state_service_pb2, task_definition_pb2
+from flyte.models import GroupData
 
 
 @dataclass
@@ -28,6 +28,7 @@ class Action:
     started: bool = False
     retries: int = 0
     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.
 
     @property
     def name(self) -> str:
@@ -91,6 +92,8 @@ class Action:
         if not self.started:
             self.task = action.task
 
+        self.cache_key = action.cache_key
+
     def set_client_error(self, exc: Exception):
         self.client_err = exc
 
@@ -106,6 +109,7 @@ class Action:
         task_spec: task_definition_pb2.TaskSpec,
         inputs_uri: str,
         run_output_base: str,
+        cache_key: str | None = None,
     ) -> Action:
         return cls(
             action_id=sub_action_id,
@@ -115,6 +119,7 @@ class Action:
             task=task_spec,
             inputs_uri=inputs_uri,
             run_output_base=run_output_base,
+            cache_key=cache_key,
         )
 
     @classmethod
@@ -130,7 +135,7 @@ class Action:
         """
         from flyte._logging import logger
 
-        logger.info(f"In Action from_state {obj.action_id} {obj.phase} {obj.output_uri}")
+        logger.debug(f"In Action from_state {obj.action_id} {obj.phase} {obj.output_uri}")
         return cls(
             action_id=obj.action_id,
             parent_action_name=parent_action_name,

flyte/_internal/controllers/remote/_client.py

@@ -20,7 +20,11 @@ class ControllerClient:
 
     @classmethod
     async def for_endpoint(cls, endpoint: str, insecure: bool = False, **kwargs) -> ControllerClient:
-        return cls(await create_channel(endpoint, insecure=insecure, **kwargs))
+        return cls(await create_channel(endpoint, None, insecure=insecure, **kwargs))
+
+    @classmethod
+    async def for_api_key(cls, api_key: str, insecure: bool = False, **kwargs) -> ControllerClient:
+        return cls(await create_channel(None, api_key, insecure=insecure, **kwargs))
 
     @property
     def state_service(self) -> StateService:

flyte/_internal/controllers/remote/_controller.py

@@ -1,9 +1,13 @@
 from __future__ import annotations
 
 import asyncio
+import concurrent.futures
+import os
+import threading
 from collections import defaultdict
+from collections.abc import Callable
 from pathlib import Path
-from typing import Any, Awaitable, DefaultDict, Tuple, TypeVar
+from typing import Any, AsyncIterable, Awaitable, DefaultDict, Tuple, TypeVar
 
 import flyte
 import flyte.errors
@@ -11,7 +15,6 @@ 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._datastructures import ActionID, NativeInterface, SerializationContext
 from flyte._internal.controllers import TraceInfo
 from flyte._internal.controllers.remote._action import Action
 from flyte._internal.controllers.remote._core import Controller
@@ -21,16 +24,18 @@ from flyte._internal.runtime.task_serde import translate_task_to_wire
 from flyte._logging import logger
 from flyte._protos.workflow import run_definition_pb2, task_definition_pb2
 from flyte._task import TaskTemplate
+from flyte._utils.helpers import _selector_policy
+from flyte.models import ActionID, NativeInterface, SerializationContext
 
 R = TypeVar("R")
 
 
-async def upload_inputs_with_retry(inputs: convert.Inputs, inputs_uri: str) -> None:
+async def upload_inputs_with_retry(serialized_inputs: AsyncIterable[bytes] | bytes, inputs_uri: str) -> None:
     """
     Upload inputs to the specified URI with error handling.
 
     Args:
-        inputs: The inputs to upload
+        serialized_inputs: The serialized inputs to upload
         inputs_uri: The destination URI
 
     Raises:
@@ -38,9 +43,9 @@ async def upload_inputs_with_retry(inputs: convert.Inputs, inputs_uri: str) -> N
     """
     try:
         # TODO Add retry decorator to this
-        await io.upload_inputs(inputs, inputs_uri)
+        await storage.put_stream(serialized_inputs, to_path=inputs_uri)
     except Exception as e:
-        logger.exception("Failed to upload inputs", e)
+        logger.exception("Failed to upload inputs")
         raise flyte.errors.RuntimeSystemError(type(e).__name__, str(e)) from e
 
 
@@ -89,6 +94,10 @@ async def load_and_convert_outputs(iface: NativeInterface, realized_outputs_uri:
     return await convert.convert_outputs_to_native(iface, outputs)
 
 
+def unique_action_name(action_id: ActionID) -> str:
+    return f"{action_id.name}_{action_id.run_name}"
+
+
 class RemoteController(Controller):
     """
     This a specialized controller that wraps the core controller and performs IO, serialization and deserialization
@@ -99,7 +108,7 @@ class RemoteController(Controller):
         client_coro: Awaitable[ClientSet],
         workers: int,
         max_system_retries: int,
-        default_parent_concurrency: int = 100,
+        default_parent_concurrency: int = 1000,
     ):
         """ """
         super().__init__(
@@ -111,31 +120,44 @@ class RemoteController(Controller):
         self._parent_action_semaphore: DefaultDict[str, asyncio.Semaphore] = defaultdict(
             lambda: asyncio.Semaphore(default_parent_concurrency)
         )
+        self._parent_action_task_call_sequence: DefaultDict[str, DefaultDict[int, int]] = defaultdict(
+            lambda: defaultdict(int)
+        )
+        self._submit_loop: asyncio.AbstractEventLoop | None = None
+        self._submit_thread: threading.Thread | None = None
 
-    async def _submit(self, _task: TaskTemplate, *args, **kwargs) -> Any:
+    def generate_task_call_sequence(self, task_obj: object, action_id: ActionID) -> int:
+        """
+        Generate a task call sequence for the given task object and action ID.
+        This is used to track the number of times a task is called within an action.
+        """
+        current_action_sequencer = self._parent_action_task_call_sequence[unique_action_name(action_id)]
+        current_task_id = id(task_obj)
+        v = current_action_sequencer[current_task_id]
+        new_seq = v + 1
+        current_action_sequencer[current_task_id] = new_seq
+        return new_seq
+
+    async def _submit(self, _task_call_seq: int, _task: TaskTemplate, *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
 
-        inputs = await convert.convert_from_native_to_inputs(_task.native_interface, *args, **kwargs)
-        sub_action_id, sub_action_output_path = convert.generate_sub_action_id_and_output_path(tctx, _task.name, inputs)
-
         # In the case of a regular code bundle, we will just pass it down as it is to the downstream tasks
         # It is not allowed to change the code bundle (for regular code bundles) in the middle of a run.
         code_bundle = tctx.code_bundle
 
         if code_bundle and code_bundle.pkl:
-            logger.debug(f"Building new pkl bundle for task {sub_action_id.name}")
+            logger.debug(f"Building new pkl bundle for task {_task.name}")
             code_bundle = await build_pkl_bundle(
                 _task,
                 upload_to_controlplane=False,
-                upload_from_dataplane_path=io.pkl_path(sub_action_output_path),
+                upload_from_dataplane_base_path=tctx.run_base_dir,
             )
 
-        inputs_uri = io.inputs_path(sub_action_output_path)
-        await upload_inputs_with_retry(inputs, inputs_uri)
+        inputs = await convert.convert_from_native_to_inputs(_task.native_interface, *args, **kwargs)
 
         root_dir = Path(code_bundle.destination).absolute() if code_bundle else Path.cwd()
         # Don't set output path in sec context because node executor will set it
@@ -146,12 +168,41 @@ class RemoteController(Controller):
             code_bundle=code_bundle,
             version=tctx.version,
             # supplied version.
-            input_path=inputs_uri,
+            # input_path=inputs_uri,
             image_cache=tctx.compiled_image_cache,
             root_dir=root_dir,
         )
 
         task_spec = translate_task_to_wire(_task, new_serialization_context)
+        serialized_inputs = inputs.proto_inputs.SerializeToString(deterministic=True)
+
+        inputs_hash = convert.generate_inputs_hash(serialized_inputs)
+        sub_action_id, sub_action_output_path = convert.generate_sub_action_id_and_output_path(
+            tctx, task_spec, inputs_hash, _task_call_seq
+        )
+
+        inputs_uri = io.inputs_path(sub_action_output_path)
+        await upload_inputs_with_retry(serialized_inputs, inputs_uri)
+
+        md = task_spec.task_template.metadata
+        ignored_input_vars = []
+        if len(md.cache_ignore_input_vars) > 0:
+            ignored_input_vars = list(md.cache_ignore_input_vars)
+        cache_key = None
+        if task_spec.task_template.metadata and task_spec.task_template.metadata.discoverable:
+            discovery_version = task_spec.task_template.metadata.discovery_version
+            cache_key = convert.generate_cache_key_hash(
+                _task.name,
+                inputs_hash,
+                task_spec.task_template.interface,
+                discovery_version,
+                ignored_input_vars,
+                inputs.proto_inputs,
+            )
+
+        # Clear to free memory
+        serialized_inputs = None  # type: ignore
+        inputs_hash = None  # type: ignore
 
         action = Action.from_task(
             sub_action_id=run_definition_pb2.ActionIdentifier(
@@ -168,6 +219,7 @@ class RemoteController(Controller):
             task_spec=task_spec,
             inputs_uri=inputs_uri,
             run_output_base=tctx.run_base_dir,
+            cache_key=cache_key,
         )
 
         try:
@@ -205,8 +257,51 @@ class RemoteController(Controller):
         if tctx is None:
             raise flyte.errors.RuntimeSystemError("BadContext", "Task context not initialized")
         current_action_id = tctx.action
-        async with self._parent_action_semaphore[current_action_id.name]:
-            return await self._submit(_task, *args, **kwargs)
+        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_call_seq, _task, *args, **kwargs)
+
+    def _sync_thread_loop_runner(self) -> None:
+        """This method runs the event loop and should be invoked in a separate thread."""
+
+        loop = self._submit_loop
+        assert loop is not None
+        try:
+            loop.run_forever()
+        finally:
+            loop.close()
+
+    def submit_sync(self, _task: TaskTemplate, *args, **kwargs) -> concurrent.futures.Future:
+        """
+        This function creates a cached thread and loop for the purpose of calling the submit method synchronously,
+        returning a concurrent Future that can be awaited. There's no need for a lock because this function itself is
+        single threaded and non-async. This pattern here is basically the trivial/degenerate case of the thread pool
+        in the LocalController.
+        Please see additional comments in protocol.
+
+        :param _task:
+        :param args:
+        :param kwargs:
+        :return:
+        """
+        if self._submit_thread is None:
+            # Please see LocalController for the general implementation of this pattern.
+            def exc_handler(loop, context):
+                logger.error(f"Remote controller submit sync loop caught exception in {loop}: {context}")
+
+            with _selector_policy():
+                self._submit_loop = asyncio.new_event_loop()
+                self._submit_loop.set_exception_handler(exc_handler)
+
+            self._submit_thread = threading.Thread(
+                name=f"remote-controller-{os.getpid()}-submitter", daemon=True, target=self._sync_thread_loop_runner
+            )
+            self._submit_thread.start()
+
+        coro = self.submit(_task, *args, **kwargs)
+        assert self._submit_loop is not None, "Submit loop should always have been initialized by now"
+        fut = asyncio.run_coroutine_threadsafe(coro, self._submit_loop)
+        return fut
 
     async def finalize_parent_action(self, action_id: ActionID):
         """
@@ -220,16 +315,17 @@ class RemoteController(Controller):
             org=action_id.org,
         )
         await super()._finalize_parent_action(run_id=run_id, parent_action_name=action_id.name)
-        self._parent_action_semaphore.pop(action_id.name, None)
+        self._parent_action_semaphore.pop(unique_action_name(action_id), None)
+        self._parent_action_task_call_sequence.pop(unique_action_name(action_id), None)
 
     async def get_action_outputs(
-        self, _interface: NativeInterface, _func_name: str, *args, **kwargs
+        self, _interface: NativeInterface, _func: Callable, *args, **kwargs
     ) -> Tuple[TraceInfo, bool]:
         """
         This method returns the outputs of the action, if it is available.
         If not available it raises a NotFoundError.
         :param _interface: NativeInterface
-        :param _func_name: Function name
+        :param _func: Function name
         :param args: Arguments
         :param kwargs: Keyword arguments
         :return:
@@ -240,11 +336,19 @@ class RemoteController(Controller):
             raise flyte.errors.RuntimeSystemError("BadContext", "Task context not initialized")
         current_action_id = tctx.action
 
+        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)
-        sub_action_id, sub_action_output_path = convert.generate_sub_action_id_and_output_path(tctx, _func_name, inputs)
+        serialized_inputs = inputs.proto_inputs.SerializeToString(deterministic=True)
+
+        sub_action_id, sub_action_output_path = convert.generate_sub_action_id_and_output_path(
+            tctx, func_name, serialized_inputs, invoke_seq_num
+        )
 
         inputs_uri = io.inputs_path(sub_action_output_path)
-        await upload_inputs_with_retry(inputs, inputs_uri)
+        await upload_inputs_with_retry(serialized_inputs, inputs_uri)
+        # Clear to free memory
+        serialized_inputs = None  # type: ignore
 
         prev_action = await self.get_action(
             run_definition_pb2.ActionIdentifier(
@@ -310,12 +414,40 @@ class RemoteController(Controller):
         current_action_id = tctx.action
         task_name = _task.spec.task_template.id.name
 
-        native_interface = types.guess_interface(_task.spec.task_template.interface)
+        invoke_seq_num = self.generate_task_call_sequence(_task, current_action_id)
+
+        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)
-        sub_action_id, sub_action_output_path = convert.generate_sub_action_id_and_output_path(tctx, task_name, inputs)
+        serialized_inputs = inputs.proto_inputs.SerializeToString(deterministic=True)
+        inputs_hash = convert.generate_inputs_hash(serialized_inputs)
+        sub_action_id, sub_action_output_path = convert.generate_sub_action_id_and_output_path(
+            tctx, task_name, inputs_hash, invoke_seq_num
+        )
 
         inputs_uri = io.inputs_path(sub_action_output_path)
-        await upload_inputs_with_retry(inputs, inputs_uri)
+        await upload_inputs_with_retry(serialized_inputs, inputs_uri)
+        # cache key - task name, task signature, inputs, cache version
+        cache_key = None
+        md = _task.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
+            cache_key = convert.generate_cache_key_hash(
+                task_name,
+                inputs_hash,
+                _task.spec.task_template.interface,
+                discovery_version,
+                ignored_input_vars,
+                inputs.proto_inputs,
+            )
+
+        # Clear to free memory
+        serialized_inputs = None  # type: ignore
+        inputs_hash = None  # type: ignore
 
         action = Action.from_task(
             sub_action_id=run_definition_pb2.ActionIdentifier(
@@ -332,6 +464,7 @@ class RemoteController(Controller):
             task_spec=_task.spec,
             inputs_uri=inputs_uri,
             run_output_base=tctx.run_base_dir,
+            cache_key=cache_key,
         )
 
         try:

flyte/_internal/controllers/remote/_core.py

@@ -7,6 +7,7 @@ from asyncio import Event
 from typing import Awaitable, Coroutine, Optional
 
 import grpc.aio
+from google.protobuf.wrappers_pb2 import StringValue
 
 import flyte.errors
 from flyte._logging import log, logger
@@ -32,7 +33,7 @@ class Controller:
         max_system_retries: int = 5,
         resource_log_interval_sec: float = 10.0,
         min_backoff_on_err_sec: float = 0.1,
-        thread_wait_timeout_sec: float = 10.0,
+        thread_wait_timeout_sec: float = 5.0,
         enqueue_timeout_sec: float = 5.0,
     ):
         """
@@ -286,10 +287,11 @@ class Controller:
         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,
-                )
+                # 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]:
@@ -310,6 +312,11 @@ class Controller:
         if not action.is_started() and action.task is not None:
             logger.debug(f"Attempting to launch action: {action.name}")
             try:
+                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)
+
                 await self._queue_service.EnqueueAction(
                     queue_service_pb2.EnqueueActionRequest(
                         action_id=action.action_id,
@@ -323,6 +330,7 @@ class Controller:
                                 name=action.task.task_template.id.name,
                             ),
                             spec=action.task,
+                            cache_key=cache_key,
                         ),
                         input_uri=action.inputs_uri,
                         run_output_base=action.run_output_base,