flyte 0.2.0b12__py3-none-any.whl → 0.2.0b14__py3-none-any.whl

flyte/__init__.py

@@ -46,10 +46,13 @@ __all__ = [
     "with_runcontext",
 ]
 
+import sys
+
 from ._cache import Cache, CachePolicy, CacheRequest
 from ._context import ctx
 from ._deploy import deploy
 from ._environment import Environment
+from ._excepthook import custom_excepthook
 from ._group import group
 from ._image import Image
 from ._initialize import init, init_from_config
@@ -64,3 +67,5 @@ from ._task_environment import TaskEnvironment
 from ._timeout import Timeout, TimeoutType
 from ._trace import trace
 from ._version import __version__
+
+sys.excepthook = custom_excepthook

flyte/_excepthook.py

@@ -0,0 +1,37 @@
+import logging
+import sys
+import traceback
+
+from flyte._logging import logger
+
+# Save the original excepthook so we can call it later
+original_excepthook = sys.excepthook
+
+# Filters: exclude frames where filename or function name contains these substrings
+EXCLUDED_MODULE_SUBSTRINGS = ["_internal", "syncify"]
+EXCLUDED_FILE_SUBSTRINGS = ["syncify", "_code_bundle"]
+
+
+def should_include_frame(frame: traceback.FrameSummary) -> bool:
+    return not (
+        any(sub in frame.name for sub in EXCLUDED_MODULE_SUBSTRINGS)
+        or any(sub in frame.filename for sub in EXCLUDED_FILE_SUBSTRINGS)
+    )
+
+
+def custom_excepthook(exc_type, exc_value, exc_tb):
+    """
+    Custom exception hook to filter and format tracebacks.
+    If the logger's level is DEBUG, it uses the original excepthook.
+    """
+
+    if logger.getEffectiveLevel() <= logging.DEBUG:
+        original_excepthook(exc_type, exc_value, exc_tb)
+    else:
+        # Extract and filter traceback
+        tb_list = traceback.extract_tb(exc_tb)
+        filtered_tb = [frame for frame in tb_list if should_include_frame(frame)]
+        # Print the filtered version (custom format)
+        print("Filtered traceback (most recent call last):")
+        print("".join(traceback.format_list(filtered_tb)))
+        print(f"{exc_type.__name__}: {exc_value}\n")

flyte/_internal/controllers/remote/_action.py

@@ -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

flyte/_internal/controllers/remote/_controller.py

@@ -174,16 +174,35 @@ class RemoteController(Controller):
         )
 
         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, serialized_inputs, _task_call_seq
+            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(
@@ -200,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:
@@ -400,14 +420,33 @@ class RemoteController(Controller):
         inputs = await convert.convert_from_native_to_inputs(native_interface, *args, **kwargs)
         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, serialized_inputs, invoke_seq_num
+            tctx, task_name, inputs_hash, invoke_seq_num
         )
 
         inputs_uri = io.inputs_path(sub_action_output_path)
         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(
@@ -424,6 +463,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
@@ -311,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,
@@ -324,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,

flyte/_internal/runtime/convert.py

@@ -2,10 +2,13 @@ from __future__ import annotations
 
 import base64
 import hashlib
+import inspect
 from dataclasses import dataclass
-from typing import Any, Dict, Tuple, Union
+from types import NoneType
+from typing import Any, Dict, List, Tuple, Union, get_args
 
 from flyteidl.core import execution_pb2, literals_pb2
+from flyteidl.core.interface_pb2 import TypedInterface
 
 import flyte.errors
 import flyte.storage as storage
@@ -57,25 +60,38 @@ async def convert_inputs_to_native(inputs: Inputs, python_interface: NativeInter
     return native_vals
 
 
+def is_optional_type(tp) -> bool:
+    """
+    True if the *annotation* `tp` is equivalent to Optional[…].
+    Works for Optional[T], Union[T, None], and T | None.
+    """
+    return NoneType in get_args(tp)  # fastest check
+
+
 async def convert_from_native_to_inputs(interface: NativeInterface, *args, **kwargs) -> Inputs:
     kwargs = interface.convert_to_kwargs(*args, **kwargs)
     if len(kwargs) == 0:
         return Inputs.empty()
+
     # fill in defaults if missing
     for input_name, (input_type, default_value) in interface.inputs.items():
         if input_name not in kwargs:
-            if default_value is not None:
+            if (default_value is not None and default_value is not inspect.Signature.empty) or (
+                default_value is None and is_optional_type(input_type)
+            ):
                 kwargs[input_name] = default_value
-            # todo: fill in Nones for optional inputs
     if len(kwargs) < len(interface.inputs):
         raise ValueError(
             f"Received {len(kwargs)} inputs but interface has {len(interface.inputs)}. "
             f"Please provide all required inputs."
         )
     literal_map = await TypeEngine.dict_to_literal_map(kwargs, interface.get_input_types())
+    # Make sure we the interface, not literal_map or kwargs, because those may have a different order
     return Inputs(
         proto_inputs=run_definition_pb2.Inputs(
-            literals=[run_definition_pb2.NamedLiteral(name=k, value=v) for k, v in literal_map.literals.items()]
+            literals=[
+                run_definition_pb2.NamedLiteral(name=k, value=literal_map.literals[k]) for k in interface.inputs.keys()
+            ]
         )
     )
 
@@ -201,10 +217,48 @@ def hash_data(data: Union[str, bytes]) -> str:
     return base64.b64encode(digest).decode("utf-8")
 
 
+def generate_inputs_hash(serialized_inputs: str | bytes) -> str:
+    """
+    Generate a hash for the inputs. This is used to uniquely identify the inputs for a task.
+    :return: A hexadecimal string representation of the hash.
+    """
+    return hash_data(serialized_inputs)
+
+
+def generate_cache_key_hash(
+    task_name: str,
+    inputs_hash: str,
+    task_interface: TypedInterface,
+    cache_version: str,
+    ignored_input_vars: List[str],
+    proto_inputs: run_definition_pb2.Inputs,
+) -> str:
+    """
+    Generate a cache key hash based on the inputs hash, task name, task interface, and cache version.
+    This is used to uniquely identify the cache key for a task.
+
+    :param task_name: The name of the task.
+    :param inputs_hash: The hash of the inputs.
+    :param task_interface: The interface of the task.
+    :param cache_version: The version of the cache.
+    :param ignored_input_vars: A list of input variable names to ignore when generating the cache key.
+    :param proto_inputs: The proto inputs for the task, only used if there are ignored inputs.
+    :return: A hexadecimal string representation of the cache key hash.
+    """
+    if ignored_input_vars:
+        filtered = [named_lit for named_lit in proto_inputs.literals if named_lit.name not in ignored_input_vars]
+        final = run_definition_pb2.Inputs(literals=filtered)
+        final_inputs = final.SerializeToString(deterministic=True)
+    else:
+        final_inputs = inputs_hash
+    data = f"{final_inputs}{task_name}{task_interface.SerializeToString(deterministic=True)}{cache_version}"
+    return hash_data(data)
+
+
 def generate_sub_action_id_and_output_path(
     tctx: TaskContext,
     task_spec_or_name: task_definition_pb2.TaskSpec | str,
-    serialized_inputs: str | bytes,
+    inputs_hash: str,
     invoke_seq: int,
 ) -> Tuple[ActionID, str]:
     """
@@ -213,14 +267,14 @@ def generate_sub_action_id_and_output_path(
     action name = current action name + task name + input hash + group name (if available)
     :param tctx:
     :param task_spec_or_name: task specification or task name. Task name is only used in case of trace actions.
-    :param serialized_inputs:
+    :param inputs_hash: Consistent hash string of the inputs
     :param invoke_seq: The sequence number of the invocation, used to differentiate between multiple invocations.
     :return:
     """
     current_action_id = tctx.action
     current_output_path = tctx.run_base_dir
-    inputs_hash = hash_data(serialized_inputs)
     if isinstance(task_spec_or_name, task_definition_pb2.TaskSpec):
+        task_spec_or_name.task_template.interface
         task_hash = hash_data(task_spec_or_name.SerializeToString(deterministic=True))
     else:
         task_hash = task_spec_or_name

flyte/_internal/runtime/task_serde.py

@@ -165,7 +165,7 @@ def get_proto_task(task: TaskTemplate, serialize_context: SerializationContext)
             discoverable=cache_enabled,
             discovery_version=cache_version,
             cache_serializable=task_cache.serialize,
-            cache_ignore_input_vars=task_cache.ignored_inputs,
+            cache_ignore_input_vars=task_cache.get_ignored_inputs() if cache_enabled else None,
             runtime=tasks_pb2.RuntimeMetadata(),
             retries=get_proto_retry_strategy(task.retries),
             timeout=get_proto_timeout(task.timeout),

flyte/_protos/workflow/queue_service_pb2.py

@@ -12,12 +12,13 @@ _sym_db = _symbol_database.Default()
 
 
 from flyteidl.core import types_pb2 as flyteidl_dot_core_dot_types__pb2
+from google.protobuf import wrappers_pb2 as google_dot_protobuf_dot_wrappers__pb2
 from flyte._protos.validate.validate import validate_pb2 as validate_dot_validate__pb2
 from flyte._protos.workflow import run_definition_pb2 as workflow_dot_run__definition__pb2
 from flyte._protos.workflow import task_definition_pb2 as workflow_dot_task__definition__pb2
 
 
-DESCRIPTOR = _descriptor_pool.Default().AddSerializedFile(b'\n\x1cworkflow/queue_service.proto\x12\x11\x63loudidl.workflow\x1a\x19\x66lyteidl/core/types.proto\x1a\x17validate/validate.proto\x1a\x1dworkflow/run_definition.proto\x1a\x1eworkflow/task_definition.proto\"\x7f\n\x10WorkerIdentifier\x12+\n\x0corganization\x18\x01 \x01(\tB\x07\xfa\x42\x04r\x02\x10\x01R\x0corganization\x12!\n\x07\x63luster\x18\x02 \x01(\tB\x07\xfa\x42\x04r\x02\x10\x01R\x07\x63luster\x12\x1b\n\x04name\x18\x03 \x01(\tB\x07\xfa\x42\x04r\x02\x10\x01R\x04name\"\x8a\x04\n\x14\x45nqueueActionRequest\x12J\n\taction_id\x18\x01 \x01(\x0b\x32#.cloudidl.workflow.ActionIdentifierB\x08\xfa\x42\x05\x8a\x01\x02\x10\x01R\x08\x61\x63tionId\x12\x31\n\x12parent_action_name\x18\x02 \x01(\tH\x01R\x10parentActionName\x88\x01\x01\x12$\n\tinput_uri\x18\x06 \x01(\tB\x07\xfa\x42\x04r\x02\x10\x01R\x08inputUri\x12/\n\x0frun_output_base\x18\x07 \x01(\tB\x07\xfa\x42\x04r\x02\x10\x01R\rrunOutputBase\x12\x14\n\x05group\x18\x08 \x01(\tR\x05group\x12\x18\n\x07subject\x18\t \x01(\tR\x07subject\x12=\n\x04task\x18\n \x01(\x0b\x32\x1d.cloudidl.workflow.TaskActionB\x08\xfa\x42\x05\x8a\x01\x02\x10\x01H\x00R\x04task\x12@\n\x05trace\x18\x0b \x01(\x0b\x32\x1e.cloudidl.workflow.TraceActionB\x08\xfa\x42\x05\x8a\x01\x02\x10\x01H\x00R\x05trace\x12L\n\tcondition\x18\x0c \x01(\x0b\x32\".cloudidl.workflow.ConditionActionB\x08\xfa\x42\x05\x8a\x01\x02\x10\x01H\x00R\tconditionB\x06\n\x04specB\x15\n\x13_parent_action_name\"z\n\nTaskAction\x12\x31\n\x02id\x18\x01 \x01(\x0b\x32!.cloudidl.workflow.TaskIdentifierR\x02id\x12\x39\n\x04spec\x18\x02 \x01(\x0b\x32\x1b.cloudidl.workflow.TaskSpecB\x08\xfa\x42\x05\x8a\x01\x02\x10\x01R\x04spec\"*\n\x0bTraceAction\x12\x1b\n\x04name\x18\x01 \x01(\tB\x07\xfa\x42\x04r\x02\x10\x01R\x04name\"\x8a\x02\n\x0f\x43onditionAction\x12\x1b\n\x04name\x18\x01 \x01(\tB\x07\xfa\x42\x04r\x02\x10\x01R\x04name\x12 \n\x06run_id\x18\x02 \x01(\tB\x07\xfa\x42\x04r\x02\x10\x01H\x00R\x05runId\x12&\n\taction_id\x18\x03 \x01(\tB\x07\xfa\x42\x04r\x02\x10\x01H\x00R\x08\x61\x63tionId\x12\x18\n\x06global\x18\x04 \x01(\x08H\x00R\x06global\x12.\n\x04type\x18\x06 \x01(\x0b\x32\x1a.flyteidl.core.LiteralTypeR\x04type\x12\x16\n\x06prompt\x18\x07 \x01(\tR\x06prompt\x12 \n\x0b\x64\x65scription\x18\x08 \x01(\tR\x0b\x64\x65scriptionB\x0c\n\x05scope\x12\x03\xf8\x42\x01\"\x17\n\x15\x45nqueueActionResponse\"Z\n\x15\x41\x62ortQueuedRunRequest\x12\x41\n\x06run_id\x18\x01 \x01(\x0b\x32 .cloudidl.workflow.RunIdentifierB\x08\xfa\x42\x05\x8a\x01\x02\x10\x01R\x05runId\"\x18\n\x16\x41\x62ortQueuedRunResponse\"\x86\x03\n\x10HeartbeatRequest\x12J\n\tworker_id\x18\x01 \x01(\x0b\x32#.cloudidl.workflow.WorkerIdentifierB\x08\xfa\x42\x05\x8a\x01\x02\x10\x01R\x08workerId\x12O\n\x11\x61\x63tive_action_ids\x18\x02 \x03(\x0b\x32#.cloudidl.workflow.ActionIdentifierR\x0f\x61\x63tiveActionIds\x12S\n\x13terminal_action_ids\x18\x03 \x03(\x0b\x32#.cloudidl.workflow.ActionIdentifierR\x11terminalActionIds\x12Q\n\x12\x61\x62orted_action_ids\x18\x04 \x03(\x0b\x32#.cloudidl.workflow.ActionIdentifierR\x10\x61\x62ortedActionIds\x12-\n\x12\x61vailable_capacity\x18\x05 \x01(\x05R\x11\x61vailableCapacity\"\xe4\x01\n\x11HeartbeatResponse\x12\x37\n\nnew_leases\x18\x01 \x03(\x0b\x32\x18.cloudidl.workflow.LeaseR\tnewLeases\x12?\n\x0e\x61\x62orted_leases\x18\x02 \x03(\x0b\x32\x18.cloudidl.workflow.LeaseR\rabortedLeases\x12U\n\x14\x66inalized_action_ids\x18\x03 \x03(\x0b\x32#.cloudidl.workflow.ActionIdentifierR\x12\x66inalizedActionIds\"a\n\x13StreamLeasesRequest\x12J\n\tworker_id\x18\x01 \x01(\x0b\x32#.cloudidl.workflow.WorkerIdentifierB\x08\xfa\x42\x05\x8a\x01\x02\x10\x01R\x08workerId\"H\n\x14StreamLeasesResponse\x12\x30\n\x06leases\x18\x01 \x03(\x0b\x32\x18.cloudidl.workflow.LeaseR\x06leases\"\xbe\x03\n\x05Lease\x12J\n\taction_id\x18\x01 \x01(\x0b\x32#.cloudidl.workflow.ActionIdentifierB\x08\xfa\x42\x05\x8a\x01\x02\x10\x01R\x08\x61\x63tionId\x12\x31\n\x12parent_action_name\x18\x02 \x01(\tH\x01R\x10parentActionName\x88\x01\x01\x12$\n\tinput_uri\x18\x04 \x01(\tB\x07\xfa\x42\x04r\x02\x10\x01R\x08inputUri\x12/\n\x0frun_output_base\x18\x05 \x01(\tB\x07\xfa\x42\x04r\x02\x10\x01R\rrunOutputBase\x12=\n\x04task\x18\x06 \x01(\x0b\x32\x1d.cloudidl.workflow.TaskActionB\x08\xfa\x42\x05\x8a\x01\x02\x10\x01H\x00R\x04task\x12L\n\tcondition\x18\x07 \x01(\x0b\x32\".cloudidl.workflow.ConditionActionB\x08\xfa\x42\x05\x8a\x01\x02\x10\x01H\x00R\tcondition\x12\x14\n\x05group\x18\x08 \x01(\tR\x05group\x12\x18\n\x07subject\x18\t \x01(\tR\x07subjectB\x0b\n\x04spec\x12\x03\xf8\x42\x01\x42\x15\n\x13_parent_action_name2\xa0\x03\n\x0cQueueService\x12\x64\n\rEnqueueAction\x12\'.cloudidl.workflow.EnqueueActionRequest\x1a(.cloudidl.workflow.EnqueueActionResponse\"\x00\x12g\n\x0e\x41\x62ortQueuedRun\x12(.cloudidl.workflow.AbortQueuedRunRequest\x1a).cloudidl.workflow.AbortQueuedRunResponse\"\x00\x12\\\n\tHeartbeat\x12#.cloudidl.workflow.HeartbeatRequest\x1a$.cloudidl.workflow.HeartbeatResponse\"\x00(\x01\x30\x01\x12\x63\n\x0cStreamLeases\x12&.cloudidl.workflow.StreamLeasesRequest\x1a\'.cloudidl.workflow.StreamLeasesResponse\"\x00\x30\x01\x42\xbe\x01\n\x15\x63om.cloudidl.workflowB\x11QueueServiceProtoH\x02P\x01Z+github.com/unionai/cloud/gen/pb-go/workflow\xa2\x02\x03\x43WX\xaa\x02\x11\x43loudidl.Workflow\xca\x02\x11\x43loudidl\\Workflow\xe2\x02\x1d\x43loudidl\\Workflow\\GPBMetadata\xea\x02\x12\x43loudidl::Workflowb\x06proto3')
+DESCRIPTOR = _descriptor_pool.Default().AddSerializedFile(b'\n\x1cworkflow/queue_service.proto\x12\x11\x63loudidl.workflow\x1a\x19\x66lyteidl/core/types.proto\x1a\x1egoogle/protobuf/wrappers.proto\x1a\x17validate/validate.proto\x1a\x1dworkflow/run_definition.proto\x1a\x1eworkflow/task_definition.proto\"\x7f\n\x10WorkerIdentifier\x12+\n\x0corganization\x18\x01 \x01(\tB\x07\xfa\x42\x04r\x02\x10\x01R\x0corganization\x12!\n\x07\x63luster\x18\x02 \x01(\tB\x07\xfa\x42\x04r\x02\x10\x01R\x07\x63luster\x12\x1b\n\x04name\x18\x03 \x01(\tB\x07\xfa\x42\x04r\x02\x10\x01R\x04name\"\x8a\x04\n\x14\x45nqueueActionRequest\x12J\n\taction_id\x18\x01 \x01(\x0b\x32#.cloudidl.workflow.ActionIdentifierB\x08\xfa\x42\x05\x8a\x01\x02\x10\x01R\x08\x61\x63tionId\x12\x31\n\x12parent_action_name\x18\x02 \x01(\tH\x01R\x10parentActionName\x88\x01\x01\x12$\n\tinput_uri\x18\x06 \x01(\tB\x07\xfa\x42\x04r\x02\x10\x01R\x08inputUri\x12/\n\x0frun_output_base\x18\x07 \x01(\tB\x07\xfa\x42\x04r\x02\x10\x01R\rrunOutputBase\x12\x14\n\x05group\x18\x08 \x01(\tR\x05group\x12\x18\n\x07subject\x18\t \x01(\tR\x07subject\x12=\n\x04task\x18\n \x01(\x0b\x32\x1d.cloudidl.workflow.TaskActionB\x08\xfa\x42\x05\x8a\x01\x02\x10\x01H\x00R\x04task\x12@\n\x05trace\x18\x0b \x01(\x0b\x32\x1e.cloudidl.workflow.TraceActionB\x08\xfa\x42\x05\x8a\x01\x02\x10\x01H\x00R\x05trace\x12L\n\tcondition\x18\x0c \x01(\x0b\x32\".cloudidl.workflow.ConditionActionB\x08\xfa\x42\x05\x8a\x01\x02\x10\x01H\x00R\tconditionB\x06\n\x04specB\x15\n\x13_parent_action_name\"\xb5\x01\n\nTaskAction\x12\x31\n\x02id\x18\x01 \x01(\x0b\x32!.cloudidl.workflow.TaskIdentifierR\x02id\x12\x39\n\x04spec\x18\x02 \x01(\x0b\x32\x1b.cloudidl.workflow.TaskSpecB\x08\xfa\x42\x05\x8a\x01\x02\x10\x01R\x04spec\x12\x39\n\tcache_key\x18\x03 \x01(\x0b\x32\x1c.google.protobuf.StringValueR\x08\x63\x61\x63heKey\"*\n\x0bTraceAction\x12\x1b\n\x04name\x18\x01 \x01(\tB\x07\xfa\x42\x04r\x02\x10\x01R\x04name\"\x8a\x02\n\x0f\x43onditionAction\x12\x1b\n\x04name\x18\x01 \x01(\tB\x07\xfa\x42\x04r\x02\x10\x01R\x04name\x12 \n\x06run_id\x18\x02 \x01(\tB\x07\xfa\x42\x04r\x02\x10\x01H\x00R\x05runId\x12&\n\taction_id\x18\x03 \x01(\tB\x07\xfa\x42\x04r\x02\x10\x01H\x00R\x08\x61\x63tionId\x12\x18\n\x06global\x18\x04 \x01(\x08H\x00R\x06global\x12.\n\x04type\x18\x06 \x01(\x0b\x32\x1a.flyteidl.core.LiteralTypeR\x04type\x12\x16\n\x06prompt\x18\x07 \x01(\tR\x06prompt\x12 \n\x0b\x64\x65scription\x18\x08 \x01(\tR\x0b\x64\x65scriptionB\x0c\n\x05scope\x12\x03\xf8\x42\x01\"\x17\n\x15\x45nqueueActionResponse\"Z\n\x15\x41\x62ortQueuedRunRequest\x12\x41\n\x06run_id\x18\x01 \x01(\x0b\x32 .cloudidl.workflow.RunIdentifierB\x08\xfa\x42\x05\x8a\x01\x02\x10\x01R\x05runId\"\x18\n\x16\x41\x62ortQueuedRunResponse\"\x86\x03\n\x10HeartbeatRequest\x12J\n\tworker_id\x18\x01 \x01(\x0b\x32#.cloudidl.workflow.WorkerIdentifierB\x08\xfa\x42\x05\x8a\x01\x02\x10\x01R\x08workerId\x12O\n\x11\x61\x63tive_action_ids\x18\x02 \x03(\x0b\x32#.cloudidl.workflow.ActionIdentifierR\x0f\x61\x63tiveActionIds\x12S\n\x13terminal_action_ids\x18\x03 \x03(\x0b\x32#.cloudidl.workflow.ActionIdentifierR\x11terminalActionIds\x12Q\n\x12\x61\x62orted_action_ids\x18\x04 \x03(\x0b\x32#.cloudidl.workflow.ActionIdentifierR\x10\x61\x62ortedActionIds\x12-\n\x12\x61vailable_capacity\x18\x05 \x01(\x05R\x11\x61vailableCapacity\"\xe4\x01\n\x11HeartbeatResponse\x12\x37\n\nnew_leases\x18\x01 \x03(\x0b\x32\x18.cloudidl.workflow.LeaseR\tnewLeases\x12?\n\x0e\x61\x62orted_leases\x18\x02 \x03(\x0b\x32\x18.cloudidl.workflow.LeaseR\rabortedLeases\x12U\n\x14\x66inalized_action_ids\x18\x03 \x03(\x0b\x32#.cloudidl.workflow.ActionIdentifierR\x12\x66inalizedActionIds\"a\n\x13StreamLeasesRequest\x12J\n\tworker_id\x18\x01 \x01(\x0b\x32#.cloudidl.workflow.WorkerIdentifierB\x08\xfa\x42\x05\x8a\x01\x02\x10\x01R\x08workerId\"H\n\x14StreamLeasesResponse\x12\x30\n\x06leases\x18\x01 \x03(\x0b\x32\x18.cloudidl.workflow.LeaseR\x06leases\"\xbe\x03\n\x05Lease\x12J\n\taction_id\x18\x01 \x01(\x0b\x32#.cloudidl.workflow.ActionIdentifierB\x08\xfa\x42\x05\x8a\x01\x02\x10\x01R\x08\x61\x63tionId\x12\x31\n\x12parent_action_name\x18\x02 \x01(\tH\x01R\x10parentActionName\x88\x01\x01\x12$\n\tinput_uri\x18\x04 \x01(\tB\x07\xfa\x42\x04r\x02\x10\x01R\x08inputUri\x12/\n\x0frun_output_base\x18\x05 \x01(\tB\x07\xfa\x42\x04r\x02\x10\x01R\rrunOutputBase\x12=\n\x04task\x18\x06 \x01(\x0b\x32\x1d.cloudidl.workflow.TaskActionB\x08\xfa\x42\x05\x8a\x01\x02\x10\x01H\x00R\x04task\x12L\n\tcondition\x18\x07 \x01(\x0b\x32\".cloudidl.workflow.ConditionActionB\x08\xfa\x42\x05\x8a\x01\x02\x10\x01H\x00R\tcondition\x12\x14\n\x05group\x18\x08 \x01(\tR\x05group\x12\x18\n\x07subject\x18\t \x01(\tR\x07subjectB\x0b\n\x04spec\x12\x03\xf8\x42\x01\x42\x15\n\x13_parent_action_name2\xa0\x03\n\x0cQueueService\x12\x64\n\rEnqueueAction\x12\'.cloudidl.workflow.EnqueueActionRequest\x1a(.cloudidl.workflow.EnqueueActionResponse\"\x00\x12g\n\x0e\x41\x62ortQueuedRun\x12(.cloudidl.workflow.AbortQueuedRunRequest\x1a).cloudidl.workflow.AbortQueuedRunResponse\"\x00\x12\\\n\tHeartbeat\x12#.cloudidl.workflow.HeartbeatRequest\x1a$.cloudidl.workflow.HeartbeatResponse\"\x00(\x01\x30\x01\x12\x63\n\x0cStreamLeases\x12&.cloudidl.workflow.StreamLeasesRequest\x1a\'.cloudidl.workflow.StreamLeasesResponse\"\x00\x30\x01\x42\xbe\x01\n\x15\x63om.cloudidl.workflowB\x11QueueServiceProtoH\x02P\x01Z+github.com/unionai/cloud/gen/pb-go/workflow\xa2\x02\x03\x43WX\xaa\x02\x11\x43loudidl.Workflow\xca\x02\x11\x43loudidl\\Workflow\xe2\x02\x1d\x43loudidl\\Workflow\\GPBMetadata\xea\x02\x12\x43loudidl::Workflowb\x06proto3')
 
 _globals = globals()
 _builder.BuildMessageAndEnumDescriptors(DESCRIPTOR, _globals)
@@ -73,32 +74,32 @@ if _descriptor._USE_C_DESCRIPTORS == False:
   _LEASE.fields_by_name['task']._serialized_options = b'\372B\005\212\001\002\020\001'
   _LEASE.fields_by_name['condition']._options = None
   _LEASE.fields_by_name['condition']._serialized_options = b'\372B\005\212\001\002\020\001'
-  _globals['_WORKERIDENTIFIER']._serialized_start=166
-  _globals['_WORKERIDENTIFIER']._serialized_end=293
-  _globals['_ENQUEUEACTIONREQUEST']._serialized_start=296
-  _globals['_ENQUEUEACTIONREQUEST']._serialized_end=818
-  _globals['_TASKACTION']._serialized_start=820
-  _globals['_TASKACTION']._serialized_end=942
-  _globals['_TRACEACTION']._serialized_start=944
-  _globals['_TRACEACTION']._serialized_end=986
-  _globals['_CONDITIONACTION']._serialized_start=989
-  _globals['_CONDITIONACTION']._serialized_end=1255
-  _globals['_ENQUEUEACTIONRESPONSE']._serialized_start=1257
-  _globals['_ENQUEUEACTIONRESPONSE']._serialized_end=1280
-  _globals['_ABORTQUEUEDRUNREQUEST']._serialized_start=1282
-  _globals['_ABORTQUEUEDRUNREQUEST']._serialized_end=1372
-  _globals['_ABORTQUEUEDRUNRESPONSE']._serialized_start=1374
-  _globals['_ABORTQUEUEDRUNRESPONSE']._serialized_end=1398
-  _globals['_HEARTBEATREQUEST']._serialized_start=1401
-  _globals['_HEARTBEATREQUEST']._serialized_end=1791
-  _globals['_HEARTBEATRESPONSE']._serialized_start=1794
-  _globals['_HEARTBEATRESPONSE']._serialized_end=2022
-  _globals['_STREAMLEASESREQUEST']._serialized_start=2024
-  _globals['_STREAMLEASESREQUEST']._serialized_end=2121
-  _globals['_STREAMLEASESRESPONSE']._serialized_start=2123
-  _globals['_STREAMLEASESRESPONSE']._serialized_end=2195
-  _globals['_LEASE']._serialized_start=2198
-  _globals['_LEASE']._serialized_end=2644
-  _globals['_QUEUESERVICE']._serialized_start=2647
-  _globals['_QUEUESERVICE']._serialized_end=3063
+  _globals['_WORKERIDENTIFIER']._serialized_start=198
+  _globals['_WORKERIDENTIFIER']._serialized_end=325
+  _globals['_ENQUEUEACTIONREQUEST']._serialized_start=328
+  _globals['_ENQUEUEACTIONREQUEST']._serialized_end=850
+  _globals['_TASKACTION']._serialized_start=853
+  _globals['_TASKACTION']._serialized_end=1034
+  _globals['_TRACEACTION']._serialized_start=1036
+  _globals['_TRACEACTION']._serialized_end=1078
+  _globals['_CONDITIONACTION']._serialized_start=1081
+  _globals['_CONDITIONACTION']._serialized_end=1347
+  _globals['_ENQUEUEACTIONRESPONSE']._serialized_start=1349
+  _globals['_ENQUEUEACTIONRESPONSE']._serialized_end=1372
+  _globals['_ABORTQUEUEDRUNREQUEST']._serialized_start=1374
+  _globals['_ABORTQUEUEDRUNREQUEST']._serialized_end=1464
+  _globals['_ABORTQUEUEDRUNRESPONSE']._serialized_start=1466
+  _globals['_ABORTQUEUEDRUNRESPONSE']._serialized_end=1490
+  _globals['_HEARTBEATREQUEST']._serialized_start=1493
+  _globals['_HEARTBEATREQUEST']._serialized_end=1883
+  _globals['_HEARTBEATRESPONSE']._serialized_start=1886
+  _globals['_HEARTBEATRESPONSE']._serialized_end=2114
+  _globals['_STREAMLEASESREQUEST']._serialized_start=2116
+  _globals['_STREAMLEASESREQUEST']._serialized_end=2213
+  _globals['_STREAMLEASESRESPONSE']._serialized_start=2215
+  _globals['_STREAMLEASESRESPONSE']._serialized_end=2287
+  _globals['_LEASE']._serialized_start=2290
+  _globals['_LEASE']._serialized_end=2736
+  _globals['_QUEUESERVICE']._serialized_start=2739
+  _globals['_QUEUESERVICE']._serialized_end=3155
 # @@protoc_insertion_point(module_scope)

flyte/_protos/workflow/queue_service_pb2.pyi

@@ -1,4 +1,5 @@
 from flyteidl.core import types_pb2 as _types_pb2
+from google.protobuf import wrappers_pb2 as _wrappers_pb2
 from flyte._protos.validate.validate import validate_pb2 as _validate_pb2
 from flyte._protos.workflow import run_definition_pb2 as _run_definition_pb2
 from flyte._protos.workflow import task_definition_pb2 as _task_definition_pb2
@@ -42,12 +43,14 @@ class EnqueueActionRequest(_message.Message):
     def __init__(self, action_id: _Optional[_Union[_run_definition_pb2.ActionIdentifier, _Mapping]] = ..., parent_action_name: _Optional[str] = ..., input_uri: _Optional[str] = ..., run_output_base: _Optional[str] = ..., group: _Optional[str] = ..., subject: _Optional[str] = ..., task: _Optional[_Union[TaskAction, _Mapping]] = ..., trace: _Optional[_Union[TraceAction, _Mapping]] = ..., condition: _Optional[_Union[ConditionAction, _Mapping]] = ...) -> None: ...
 
 class TaskAction(_message.Message):
-    __slots__ = ["id", "spec"]
+    __slots__ = ["id", "spec", "cache_key"]
     ID_FIELD_NUMBER: _ClassVar[int]
     SPEC_FIELD_NUMBER: _ClassVar[int]
+    CACHE_KEY_FIELD_NUMBER: _ClassVar[int]
     id: _task_definition_pb2.TaskIdentifier
     spec: _task_definition_pb2.TaskSpec
-    def __init__(self, id: _Optional[_Union[_task_definition_pb2.TaskIdentifier, _Mapping]] = ..., spec: _Optional[_Union[_task_definition_pb2.TaskSpec, _Mapping]] = ...) -> None: ...
+    cache_key: _wrappers_pb2.StringValue
+    def __init__(self, id: _Optional[_Union[_task_definition_pb2.TaskIdentifier, _Mapping]] = ..., spec: _Optional[_Union[_task_definition_pb2.TaskSpec, _Mapping]] = ..., cache_key: _Optional[_Union[_wrappers_pb2.StringValue, _Mapping]] = ...) -> None: ...
 
 class TraceAction(_message.Message):
     __slots__ = ["name"]

flyte/_protos/workflow/state_service_pb2.py

@@ -12,12 +12,12 @@ _sym_db = _symbol_database.Default()
 
 
 from flyteidl.core import execution_pb2 as flyteidl_dot_core_dot_execution__pb2
+from google.rpc import status_pb2 as google_dot_rpc_dot_status__pb2
 from flyte._protos.validate.validate import validate_pb2 as validate_dot_validate__pb2
 from flyte._protos.workflow import run_definition_pb2 as workflow_dot_run__definition__pb2
-from flyte._protos.workflow import task_definition_pb2 as workflow_dot_task__definition__pb2
 
 
-DESCRIPTOR = _descriptor_pool.Default().AddSerializedFile(b'\n\x1cworkflow/state_service.proto\x12\x11\x63loudidl.workflow\x1a\x1d\x66lyteidl/core/execution.proto\x1a\x17validate/validate.proto\x1a\x1dworkflow/run_definition.proto\x1a\x1eworkflow/task_definition.proto\"\xc3\x01\n\x0cStoreRequest\x12J\n\taction_id\x18\x01 \x01(\x0b\x32#.cloudidl.workflow.ActionIdentifierB\x08\xfa\x42\x05\x8a\x01\x02\x10\x01R\x08\x61\x63tionId\x12\x31\n\x12parent_action_name\x18\x02 \x01(\tH\x00R\x10parentActionName\x88\x01\x01\x12\x1d\n\x05state\x18\x03 \x01(\tB\x07\xfa\x42\x04r\x02\x10\x01R\x05stateB\x15\n\x13_parent_action_name\"\x0f\n\rStoreResponse\"\xb6\x01\n\x0bLoadRequest\x12J\n\taction_id\x18\x01 \x01(\x0b\x32#.cloudidl.workflow.ActionIdentifierB\x08\xfa\x42\x05\x8a\x01\x02\x10\x01R\x08\x61\x63tionId\x12:\n\x07task_id\x18\x02 \x01(\x0b\x32!.cloudidl.workflow.TaskIdentifierR\x06taskId\x12\x1f\n\x0binputs_hash\x18\x03 \x01(\x03R\ninputsHash\"-\n\x0cLoadResponse\x12\x1d\n\x05state\x18\x01 \x01(\tB\x07\xfa\x42\x04r\x02\x10\x01R\x05state\"n\n\x0cWatchRequest\x12O\n\x10parent_action_id\x18\x01 \x01(\x0b\x32#.cloudidl.workflow.ActionIdentifierH\x00R\x0eparentActionIdB\r\n\x06\x66ilter\x12\x03\xf8\x42\x01\"\xb0\x01\n\rWatchResponse\x12\x46\n\raction_update\x18\x01 \x01(\x0b\x32\x1f.cloudidl.workflow.ActionUpdateH\x00R\x0c\x61\x63tionUpdate\x12L\n\x0f\x63ontrol_message\x18\x02 \x01(\x0b\x32!.cloudidl.workflow.ControlMessageH\x00R\x0e\x63ontrolMessageB\t\n\x07message\",\n\x0e\x43ontrolMessage\x12\x1a\n\x08sentinel\x18\x01 \x01(\x08R\x08sentinel\"\xed\x01\n\x0c\x41\x63tionUpdate\x12J\n\taction_id\x18\x01 \x01(\x0b\x32#.cloudidl.workflow.ActionIdentifierB\x08\xfa\x42\x05\x8a\x01\x02\x10\x01R\x08\x61\x63tionId\x12.\n\x05phase\x18\x02 \x01(\x0e\x32\x18.cloudidl.workflow.PhaseR\x05phase\x12\x38\n\x05\x65rror\x18\x03 \x01(\x0b\x32\x1d.flyteidl.core.ExecutionErrorH\x00R\x05\x65rror\x88\x01\x01\x12\x1d\n\noutput_uri\x18\x04 \x01(\tR\toutputUriB\x08\n\x06_error2\xf7\x01\n\x0cStateService\x12L\n\x05Store\x12\x1f.cloudidl.workflow.StoreRequest\x1a .cloudidl.workflow.StoreResponse\"\x00\x12I\n\x04Load\x12\x1e.cloudidl.workflow.LoadRequest\x1a\x1f.cloudidl.workflow.LoadResponse\"\x00\x12N\n\x05Watch\x12\x1f.cloudidl.workflow.WatchRequest\x1a .cloudidl.workflow.WatchResponse\"\x00\x30\x01\x42\xbe\x01\n\x15\x63om.cloudidl.workflowB\x11StateServiceProtoH\x02P\x01Z+github.com/unionai/cloud/gen/pb-go/workflow\xa2\x02\x03\x43WX\xaa\x02\x11\x43loudidl.Workflow\xca\x02\x11\x43loudidl\\Workflow\xe2\x02\x1d\x43loudidl\\Workflow\\GPBMetadata\xea\x02\x12\x43loudidl::Workflowb\x06proto3')
+DESCRIPTOR = _descriptor_pool.Default().AddSerializedFile(b'\n\x1cworkflow/state_service.proto\x12\x11\x63loudidl.workflow\x1a\x1d\x66lyteidl/core/execution.proto\x1a\x17google/rpc/status.proto\x1a\x17validate/validate.proto\x1a\x1dworkflow/run_definition.proto\"\xc1\x01\n\nPutRequest\x12J\n\taction_id\x18\x01 \x01(\x0b\x32#.cloudidl.workflow.ActionIdentifierB\x08\xfa\x42\x05\x8a\x01\x02\x10\x01R\x08\x61\x63tionId\x12\x31\n\x12parent_action_name\x18\x02 \x01(\tH\x00R\x10parentActionName\x88\x01\x01\x12\x1d\n\x05state\x18\x03 \x01(\tB\x07\xfa\x42\x04r\x02\x10\x01R\x05stateB\x15\n\x13_parent_action_name\"\x8f\x01\n\x0bPutResponse\x12J\n\taction_id\x18\x01 \x01(\x0b\x32#.cloudidl.workflow.ActionIdentifierB\x08\xfa\x42\x05\x8a\x01\x02\x10\x01R\x08\x61\x63tionId\x12\x34\n\x06status\x18\x02 \x01(\x0b\x32\x12.google.rpc.StatusB\x08\xfa\x42\x05\x8a\x01\x02\x10\x01R\x06status\"X\n\nGetRequest\x12J\n\taction_id\x18\x01 \x01(\x0b\x32#.cloudidl.workflow.ActionIdentifierB\x08\xfa\x42\x05\x8a\x01\x02\x10\x01R\x08\x61\x63tionId\"\xae\x01\n\x0bGetResponse\x12J\n\taction_id\x18\x01 \x01(\x0b\x32#.cloudidl.workflow.ActionIdentifierB\x08\xfa\x42\x05\x8a\x01\x02\x10\x01R\x08\x61\x63tionId\x12\x34\n\x06status\x18\x02 \x01(\x0b\x32\x12.google.rpc.StatusB\x08\xfa\x42\x05\x8a\x01\x02\x10\x01R\x06status\x12\x1d\n\x05state\x18\x03 \x01(\tB\x07\xfa\x42\x04r\x02\x10\x01R\x05state\"n\n\x0cWatchRequest\x12O\n\x10parent_action_id\x18\x01 \x01(\x0b\x32#.cloudidl.workflow.ActionIdentifierH\x00R\x0eparentActionIdB\r\n\x06\x66ilter\x12\x03\xf8\x42\x01\"\xb0\x01\n\rWatchResponse\x12\x46\n\raction_update\x18\x01 \x01(\x0b\x32\x1f.cloudidl.workflow.ActionUpdateH\x00R\x0c\x61\x63tionUpdate\x12L\n\x0f\x63ontrol_message\x18\x02 \x01(\x0b\x32!.cloudidl.workflow.ControlMessageH\x00R\x0e\x63ontrolMessageB\t\n\x07message\",\n\x0e\x43ontrolMessage\x12\x1a\n\x08sentinel\x18\x01 \x01(\x08R\x08sentinel\"\xed\x01\n\x0c\x41\x63tionUpdate\x12J\n\taction_id\x18\x01 \x01(\x0b\x32#.cloudidl.workflow.ActionIdentifierB\x08\xfa\x42\x05\x8a\x01\x02\x10\x01R\x08\x61\x63tionId\x12.\n\x05phase\x18\x02 \x01(\x0e\x32\x18.cloudidl.workflow.PhaseR\x05phase\x12\x38\n\x05\x65rror\x18\x03 \x01(\x0b\x32\x1d.flyteidl.core.ExecutionErrorH\x00R\x05\x65rror\x88\x01\x01\x12\x1d\n\noutput_uri\x18\x04 \x01(\tR\toutputUriB\x08\n\x06_error2\xf6\x01\n\x0cStateService\x12J\n\x03Put\x12\x1d.cloudidl.workflow.PutRequest\x1a\x1e.cloudidl.workflow.PutResponse\"\x00(\x01\x30\x01\x12J\n\x03Get\x12\x1d.cloudidl.workflow.GetRequest\x1a\x1e.cloudidl.workflow.GetResponse\"\x00(\x01\x30\x01\x12N\n\x05Watch\x12\x1f.cloudidl.workflow.WatchRequest\x1a .cloudidl.workflow.WatchResponse\"\x00\x30\x01\x42\xbe\x01\n\x15\x63om.cloudidl.workflowB\x11StateServiceProtoH\x02P\x01Z+github.com/unionai/cloud/gen/pb-go/workflow\xa2\x02\x03\x43WX\xaa\x02\x11\x43loudidl.Workflow\xca\x02\x11\x43loudidl\\Workflow\xe2\x02\x1d\x43loudidl\\Workflow\\GPBMetadata\xea\x02\x12\x43loudidl::Workflowb\x06proto3')
 
 _globals = globals()
 _builder.BuildMessageAndEnumDescriptors(DESCRIPTOR, _globals)
@@ -25,34 +25,42 @@ _builder.BuildTopDescriptorsAndMessages(DESCRIPTOR, 'workflow.state_service_pb2'
 if _descriptor._USE_C_DESCRIPTORS == False:
   DESCRIPTOR._options = None
   DESCRIPTOR._serialized_options = b'\n\025com.cloudidl.workflowB\021StateServiceProtoH\002P\001Z+github.com/unionai/cloud/gen/pb-go/workflow\242\002\003CWX\252\002\021Cloudidl.Workflow\312\002\021Cloudidl\\Workflow\342\002\035Cloudidl\\Workflow\\GPBMetadata\352\002\022Cloudidl::Workflow'
-  _STOREREQUEST.fields_by_name['action_id']._options = None
-  _STOREREQUEST.fields_by_name['action_id']._serialized_options = b'\372B\005\212\001\002\020\001'
-  _STOREREQUEST.fields_by_name['state']._options = None
-  _STOREREQUEST.fields_by_name['state']._serialized_options = b'\372B\004r\002\020\001'
-  _LOADREQUEST.fields_by_name['action_id']._options = None
-  _LOADREQUEST.fields_by_name['action_id']._serialized_options = b'\372B\005\212\001\002\020\001'
-  _LOADRESPONSE.fields_by_name['state']._options = None
-  _LOADRESPONSE.fields_by_name['state']._serialized_options = b'\372B\004r\002\020\001'
+  _PUTREQUEST.fields_by_name['action_id']._options = None
+  _PUTREQUEST.fields_by_name['action_id']._serialized_options = b'\372B\005\212\001\002\020\001'
+  _PUTREQUEST.fields_by_name['state']._options = None
+  _PUTREQUEST.fields_by_name['state']._serialized_options = b'\372B\004r\002\020\001'
+  _PUTRESPONSE.fields_by_name['action_id']._options = None
+  _PUTRESPONSE.fields_by_name['action_id']._serialized_options = b'\372B\005\212\001\002\020\001'
+  _PUTRESPONSE.fields_by_name['status']._options = None
+  _PUTRESPONSE.fields_by_name['status']._serialized_options = b'\372B\005\212\001\002\020\001'
+  _GETREQUEST.fields_by_name['action_id']._options = None
+  _GETREQUEST.fields_by_name['action_id']._serialized_options = b'\372B\005\212\001\002\020\001'
+  _GETRESPONSE.fields_by_name['action_id']._options = None
+  _GETRESPONSE.fields_by_name['action_id']._serialized_options = b'\372B\005\212\001\002\020\001'
+  _GETRESPONSE.fields_by_name['status']._options = None
+  _GETRESPONSE.fields_by_name['status']._serialized_options = b'\372B\005\212\001\002\020\001'
+  _GETRESPONSE.fields_by_name['state']._options = None
+  _GETRESPONSE.fields_by_name['state']._serialized_options = b'\372B\004r\002\020\001'
   _WATCHREQUEST.oneofs_by_name['filter']._options = None
   _WATCHREQUEST.oneofs_by_name['filter']._serialized_options = b'\370B\001'
   _ACTIONUPDATE.fields_by_name['action_id']._options = None
   _ACTIONUPDATE.fields_by_name['action_id']._serialized_options = b'\372B\005\212\001\002\020\001'
-  _globals['_STOREREQUEST']._serialized_start=171
-  _globals['_STOREREQUEST']._serialized_end=366
-  _globals['_STORERESPONSE']._serialized_start=368
-  _globals['_STORERESPONSE']._serialized_end=383
-  _globals['_LOADREQUEST']._serialized_start=386
-  _globals['_LOADREQUEST']._serialized_end=568
-  _globals['_LOADRESPONSE']._serialized_start=570
-  _globals['_LOADRESPONSE']._serialized_end=615
-  _globals['_WATCHREQUEST']._serialized_start=617
-  _globals['_WATCHREQUEST']._serialized_end=727
-  _globals['_WATCHRESPONSE']._serialized_start=730
-  _globals['_WATCHRESPONSE']._serialized_end=906
-  _globals['_CONTROLMESSAGE']._serialized_start=908
-  _globals['_CONTROLMESSAGE']._serialized_end=952
-  _globals['_ACTIONUPDATE']._serialized_start=955
-  _globals['_ACTIONUPDATE']._serialized_end=1192
-  _globals['_STATESERVICE']._serialized_start=1195
-  _globals['_STATESERVICE']._serialized_end=1442
+  _globals['_PUTREQUEST']._serialized_start=164
+  _globals['_PUTREQUEST']._serialized_end=357
+  _globals['_PUTRESPONSE']._serialized_start=360
+  _globals['_PUTRESPONSE']._serialized_end=503
+  _globals['_GETREQUEST']._serialized_start=505
+  _globals['_GETREQUEST']._serialized_end=593
+  _globals['_GETRESPONSE']._serialized_start=596
+  _globals['_GETRESPONSE']._serialized_end=770
+  _globals['_WATCHREQUEST']._serialized_start=772
+  _globals['_WATCHREQUEST']._serialized_end=882
+  _globals['_WATCHRESPONSE']._serialized_start=885
+  _globals['_WATCHRESPONSE']._serialized_end=1061
+  _globals['_CONTROLMESSAGE']._serialized_start=1063
+  _globals['_CONTROLMESSAGE']._serialized_end=1107
+  _globals['_ACTIONUPDATE']._serialized_start=1110
+  _globals['_ACTIONUPDATE']._serialized_end=1347
+  _globals['_STATESERVICE']._serialized_start=1350
+  _globals['_STATESERVICE']._serialized_end=1596
 # @@protoc_insertion_point(module_scope)

flyte/_protos/workflow/state_service_pb2.pyi

@@ -1,14 +1,14 @@
 from flyteidl.core import execution_pb2 as _execution_pb2
+from google.rpc import status_pb2 as _status_pb2
 from flyte._protos.validate.validate import validate_pb2 as _validate_pb2
 from flyte._protos.workflow import run_definition_pb2 as _run_definition_pb2
-from flyte._protos.workflow import task_definition_pb2 as _task_definition_pb2
 from google.protobuf import descriptor as _descriptor
 from google.protobuf import message as _message
 from typing import ClassVar as _ClassVar, Mapping as _Mapping, Optional as _Optional, Union as _Union
 
 DESCRIPTOR: _descriptor.FileDescriptor
 
-class StoreRequest(_message.Message):
+class PutRequest(_message.Message):
     __slots__ = ["action_id", "parent_action_name", "state"]
     ACTION_ID_FIELD_NUMBER: _ClassVar[int]
     PARENT_ACTION_NAME_FIELD_NUMBER: _ClassVar[int]
@@ -18,25 +18,29 @@ class StoreRequest(_message.Message):
     state: str
     def __init__(self, action_id: _Optional[_Union[_run_definition_pb2.ActionIdentifier, _Mapping]] = ..., parent_action_name: _Optional[str] = ..., state: _Optional[str] = ...) -> None: ...
 
-class StoreResponse(_message.Message):
-    __slots__ = []
-    def __init__(self) -> None: ...
+class PutResponse(_message.Message):
+    __slots__ = ["action_id", "status"]
+    ACTION_ID_FIELD_NUMBER: _ClassVar[int]
+    STATUS_FIELD_NUMBER: _ClassVar[int]
+    action_id: _run_definition_pb2.ActionIdentifier
+    status: _status_pb2.Status
+    def __init__(self, action_id: _Optional[_Union[_run_definition_pb2.ActionIdentifier, _Mapping]] = ..., status: _Optional[_Union[_status_pb2.Status, _Mapping]] = ...) -> None: ...
 
-class LoadRequest(_message.Message):
-    __slots__ = ["action_id", "task_id", "inputs_hash"]
+class GetRequest(_message.Message):
+    __slots__ = ["action_id"]
     ACTION_ID_FIELD_NUMBER: _ClassVar[int]
-    TASK_ID_FIELD_NUMBER: _ClassVar[int]
-    INPUTS_HASH_FIELD_NUMBER: _ClassVar[int]
     action_id: _run_definition_pb2.ActionIdentifier
-    task_id: _task_definition_pb2.TaskIdentifier
-    inputs_hash: int
-    def __init__(self, action_id: _Optional[_Union[_run_definition_pb2.ActionIdentifier, _Mapping]] = ..., task_id: _Optional[_Union[_task_definition_pb2.TaskIdentifier, _Mapping]] = ..., inputs_hash: _Optional[int] = ...) -> None: ...
+    def __init__(self, action_id: _Optional[_Union[_run_definition_pb2.ActionIdentifier, _Mapping]] = ...) -> None: ...
 
-class LoadResponse(_message.Message):
-    __slots__ = ["state"]
+class GetResponse(_message.Message):
+    __slots__ = ["action_id", "status", "state"]
+    ACTION_ID_FIELD_NUMBER: _ClassVar[int]
+    STATUS_FIELD_NUMBER: _ClassVar[int]
     STATE_FIELD_NUMBER: _ClassVar[int]
+    action_id: _run_definition_pb2.ActionIdentifier
+    status: _status_pb2.Status
     state: str
-    def __init__(self, state: _Optional[str] = ...) -> None: ...
+    def __init__(self, action_id: _Optional[_Union[_run_definition_pb2.ActionIdentifier, _Mapping]] = ..., status: _Optional[_Union[_status_pb2.Status, _Mapping]] = ..., state: _Optional[str] = ...) -> None: ...
 
 class WatchRequest(_message.Message):
     __slots__ = ["parent_action_id"]