flyte 2.0.0b0__py3-none-any.whl → 2.0.0b1__py3-none-any.whl

flyte/_internal/controllers/_trace.py

@@ -1,6 +1,8 @@
 from dataclasses import dataclass, field
 from typing import Any, Optional
 
+from flyteidl.core import interface_pb2
+
 from flyte.models import ActionID, NativeInterface
 
 
@@ -19,6 +21,7 @@ class TraceInfo:
     end_time: float = field(init=False, default=0.0)
     output: Optional[Any] = None
     error: Optional[Exception] = None
+    typed_interface: Optional[interface_pb2.TypedInterface] = None
 
     def add_outputs(self, output: Any, start_time: float, end_time: float):
         """

flyte/_internal/controllers/remote/_action.py

@@ -3,7 +3,7 @@ from __future__ import annotations
 from dataclasses import dataclass
 from typing import Literal
 
-from flyteidl.core import execution_pb2
+from flyteidl.core import execution_pb2, interface_pb2, tasks_pb2
 from google.protobuf import timestamp_pb2
 
 from flyte._protos.common import identifier_pb2
@@ -170,6 +170,7 @@ class Action:
         end_time: float,  # Unix timestamp in seconds with fractional seconds
         run_output_base: str,
         report_uri: str | None = None,
+        typed_interface: interface_pb2.TypedInterface | None = None,
     ) -> Action:
         """
         This creates a new action for tracing purposes. It is used to track the execution of a trace.
@@ -182,6 +183,12 @@ class Action:
         et.FromSeconds(int(end_time))
         et.nanos = int((end_time % 1) * 1e9)
 
+        spec = (
+            task_definition_pb2.TaskSpec(task_template=tasks_pb2.TaskTemplate(interface=typed_interface))
+            if typed_interface
+            else None
+        )
+
         return cls(
             action_id=action_id,
             parent_action_name=parent_action_name,
@@ -201,5 +208,6 @@ class Action:
                     output_uri=outputs_uri,
                     report_uri=report_uri,
                 ),
+                spec=spec,
             ),
         )

flyte/_internal/controllers/remote/_controller.py

@@ -21,6 +21,7 @@ from flyte._internal.controllers.remote._core import Controller
 from flyte._internal.controllers.remote._service_protocol import ClientSet
 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, task_definition_pb2
@@ -426,6 +427,9 @@ class RemoteController(Controller):
                 await io.upload_error(err.err, sub_run_output_path)
             else:
                 raise flyte.errors.RuntimeSystemError("BadTraceInfo", "Trace info does not have output or error")
+
+            typed_interface = transform_native_to_typed_interface(info.interface)
+
             trace_action = Action.from_trace(
                 parent_action_name=current_action_id.name,
                 action_id=identifier_pb2.ActionIdentifier(
@@ -444,6 +448,7 @@ class RemoteController(Controller):
                 run_output_base=tctx.run_base_dir,
                 start_time=info.start_time,
                 end_time=info.end_time,
+                typed_interface=typed_interface if typed_interface else None,
             )
             try:
                 logger.info(

flyte/_protos/workflow/queue_service_pb2.py

@@ -20,7 +20,7 @@ from flyte._protos.workflow import run_definition_pb2 as workflow_dot_run__defin
 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\x17\x63ommon/identifier.proto\x1a\x19\x66lyteidl/core/types.proto\x1a\x1fgoogle/protobuf/timestamp.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\"\xbf\x04\n\x14\x45nqueueActionRequest\x12H\n\taction_id\x18\x01 \x01(\x0b\x32!.cloudidl.common.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\x35\n\x08run_spec\x18\x03 \x01(\x0b\x32\x1a.cloudidl.workflow.RunSpecR\x07runSpec\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\"\x9d\x02\n\x0bTraceAction\x12\x1b\n\x04name\x18\x01 \x01(\tB\x07\xfa\x42\x04r\x02\x10\x01R\x04name\x12.\n\x05phase\x18\x02 \x01(\x0e\x32\x18.cloudidl.workflow.PhaseR\x05phase\x12\x39\n\nstart_time\x18\x03 \x01(\x0b\x32\x1a.google.protobuf.TimestampR\tstartTime\x12:\n\x08\x65nd_time\x18\x04 \x01(\x0b\x32\x1a.google.protobuf.TimestampH\x00R\x07\x65ndTime\x88\x01\x01\x12=\n\x07outputs\x18\x05 \x01(\x0b\x32#.cloudidl.workflow.OutputReferencesR\x07outputsB\x0b\n\t_end_time\"\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\"X\n\x15\x41\x62ortQueuedRunRequest\x12?\n\x06run_id\x18\x01 \x01(\x0b\x32\x1e.cloudidl.common.RunIdentifierB\x08\xfa\x42\x05\x8a\x01\x02\x10\x01R\x05runId\"\x18\n\x16\x41\x62ortQueuedRunResponse\"\x80\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\x12M\n\x11\x61\x63tive_action_ids\x18\x02 \x03(\x0b\x32!.cloudidl.common.ActionIdentifierR\x0f\x61\x63tiveActionIds\x12Q\n\x13terminal_action_ids\x18\x03 \x03(\x0b\x32!.cloudidl.common.ActionIdentifierR\x11terminalActionIds\x12O\n\x12\x61\x62orted_action_ids\x18\x04 \x03(\x0b\x32!.cloudidl.common.ActionIdentifierR\x10\x61\x62ortedActionIds\x12-\n\x12\x61vailable_capacity\x18\x05 \x01(\x05R\x11\x61vailableCapacity\"\xe2\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\x12S\n\x14\x66inalized_action_ids\x18\x03 \x03(\x0b\x32!.cloudidl.common.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\"\xc9\x04\n\x05Lease\x12H\n\taction_id\x18\x01 \x01(\x0b\x32!.cloudidl.common.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\x35\n\x08run_spec\x18\x03 \x01(\x0b\x32\x1a.cloudidl.workflow.RunSpecR\x07runSpec\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@\n\x05trace\x18\n \x01(\x0b\x32\x1e.cloudidl.workflow.TraceActionB\x08\xfa\x42\x05\x8a\x01\x02\x10\x01H\x00R\x05trace\x12\x14\n\x05group\x18\x08 \x01(\tR\x05group\x12\x18\n\x07subject\x18\t \x01(\tR\x07subject\x12\x12\n\x04host\x18\x0b \x01(\tR\x04hostB\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\x17\x63ommon/identifier.proto\x1a\x19\x66lyteidl/core/types.proto\x1a\x1fgoogle/protobuf/timestamp.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\"\xbf\x04\n\x14\x45nqueueActionRequest\x12H\n\taction_id\x18\x01 \x01(\x0b\x32!.cloudidl.common.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\x35\n\x08run_spec\x18\x03 \x01(\x0b\x32\x1a.cloudidl.workflow.RunSpecR\x07runSpec\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\"\xd8\x02\n\x0bTraceAction\x12\x1b\n\x04name\x18\x01 \x01(\tB\x07\xfa\x42\x04r\x02\x10\x01R\x04name\x12.\n\x05phase\x18\x02 \x01(\x0e\x32\x18.cloudidl.workflow.PhaseR\x05phase\x12\x39\n\nstart_time\x18\x03 \x01(\x0b\x32\x1a.google.protobuf.TimestampR\tstartTime\x12:\n\x08\x65nd_time\x18\x04 \x01(\x0b\x32\x1a.google.protobuf.TimestampH\x00R\x07\x65ndTime\x88\x01\x01\x12=\n\x07outputs\x18\x05 \x01(\x0b\x32#.cloudidl.workflow.OutputReferencesR\x07outputs\x12\x39\n\x04spec\x18\x06 \x01(\x0b\x32\x1b.cloudidl.workflow.TaskSpecB\x08\xfa\x42\x05\x8a\x01\x02\x10\x01R\x04specB\x0b\n\t_end_time\"\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\"X\n\x15\x41\x62ortQueuedRunRequest\x12?\n\x06run_id\x18\x01 \x01(\x0b\x32\x1e.cloudidl.common.RunIdentifierB\x08\xfa\x42\x05\x8a\x01\x02\x10\x01R\x05runId\"\x18\n\x16\x41\x62ortQueuedRunResponse\"\x80\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\x12M\n\x11\x61\x63tive_action_ids\x18\x02 \x03(\x0b\x32!.cloudidl.common.ActionIdentifierR\x0f\x61\x63tiveActionIds\x12Q\n\x13terminal_action_ids\x18\x03 \x03(\x0b\x32!.cloudidl.common.ActionIdentifierR\x11terminalActionIds\x12O\n\x12\x61\x62orted_action_ids\x18\x04 \x03(\x0b\x32!.cloudidl.common.ActionIdentifierR\x10\x61\x62ortedActionIds\x12-\n\x12\x61vailable_capacity\x18\x05 \x01(\x05R\x11\x61vailableCapacity\"\xe2\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\x12S\n\x14\x66inalized_action_ids\x18\x03 \x03(\x0b\x32!.cloudidl.common.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\"\xc9\x04\n\x05Lease\x12H\n\taction_id\x18\x01 \x01(\x0b\x32!.cloudidl.common.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\x35\n\x08run_spec\x18\x03 \x01(\x0b\x32\x1a.cloudidl.workflow.RunSpecR\x07runSpec\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@\n\x05trace\x18\n \x01(\x0b\x32\x1e.cloudidl.workflow.TraceActionB\x08\xfa\x42\x05\x8a\x01\x02\x10\x01H\x00R\x05trace\x12\x14\n\x05group\x18\x08 \x01(\tR\x05group\x12\x18\n\x07subject\x18\t \x01(\tR\x07subject\x12\x12\n\x04host\x18\x0b \x01(\tR\x04hostB\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)
@@ -50,6 +50,8 @@ if _descriptor._USE_C_DESCRIPTORS == False:
   _TASKACTION.fields_by_name['spec']._serialized_options = b'\372B\005\212\001\002\020\001'
   _TRACEACTION.fields_by_name['name']._options = None
   _TRACEACTION.fields_by_name['name']._serialized_options = b'\372B\004r\002\020\001'
+  _TRACEACTION.fields_by_name['spec']._options = None
+  _TRACEACTION.fields_by_name['spec']._serialized_options = b'\372B\005\212\001\002\020\001'
   _CONDITIONACTION.oneofs_by_name['scope']._options = None
   _CONDITIONACTION.oneofs_by_name['scope']._serialized_options = b'\370B\001'
   _CONDITIONACTION.fields_by_name['name']._options = None
@@ -85,25 +87,25 @@ if _descriptor._USE_C_DESCRIPTORS == False:
   _globals['_TASKACTION']._serialized_start=964
   _globals['_TASKACTION']._serialized_end=1145
   _globals['_TRACEACTION']._serialized_start=1148
-  _globals['_TRACEACTION']._serialized_end=1433
-  _globals['_CONDITIONACTION']._serialized_start=1436
-  _globals['_CONDITIONACTION']._serialized_end=1702
-  _globals['_ENQUEUEACTIONRESPONSE']._serialized_start=1704
-  _globals['_ENQUEUEACTIONRESPONSE']._serialized_end=1727
-  _globals['_ABORTQUEUEDRUNREQUEST']._serialized_start=1729
-  _globals['_ABORTQUEUEDRUNREQUEST']._serialized_end=1817
-  _globals['_ABORTQUEUEDRUNRESPONSE']._serialized_start=1819
-  _globals['_ABORTQUEUEDRUNRESPONSE']._serialized_end=1843
-  _globals['_HEARTBEATREQUEST']._serialized_start=1846
-  _globals['_HEARTBEATREQUEST']._serialized_end=2230
-  _globals['_HEARTBEATRESPONSE']._serialized_start=2233
-  _globals['_HEARTBEATRESPONSE']._serialized_end=2459
-  _globals['_STREAMLEASESREQUEST']._serialized_start=2461
-  _globals['_STREAMLEASESREQUEST']._serialized_end=2558
-  _globals['_STREAMLEASESRESPONSE']._serialized_start=2560
-  _globals['_STREAMLEASESRESPONSE']._serialized_end=2632
-  _globals['_LEASE']._serialized_start=2635
-  _globals['_LEASE']._serialized_end=3220
-  _globals['_QUEUESERVICE']._serialized_start=3223
-  _globals['_QUEUESERVICE']._serialized_end=3639
+  _globals['_TRACEACTION']._serialized_end=1492
+  _globals['_CONDITIONACTION']._serialized_start=1495
+  _globals['_CONDITIONACTION']._serialized_end=1761
+  _globals['_ENQUEUEACTIONRESPONSE']._serialized_start=1763
+  _globals['_ENQUEUEACTIONRESPONSE']._serialized_end=1786
+  _globals['_ABORTQUEUEDRUNREQUEST']._serialized_start=1788
+  _globals['_ABORTQUEUEDRUNREQUEST']._serialized_end=1876
+  _globals['_ABORTQUEUEDRUNRESPONSE']._serialized_start=1878
+  _globals['_ABORTQUEUEDRUNRESPONSE']._serialized_end=1902
+  _globals['_HEARTBEATREQUEST']._serialized_start=1905
+  _globals['_HEARTBEATREQUEST']._serialized_end=2289
+  _globals['_HEARTBEATRESPONSE']._serialized_start=2292
+  _globals['_HEARTBEATRESPONSE']._serialized_end=2518
+  _globals['_STREAMLEASESREQUEST']._serialized_start=2520
+  _globals['_STREAMLEASESREQUEST']._serialized_end=2617
+  _globals['_STREAMLEASESRESPONSE']._serialized_start=2619
+  _globals['_STREAMLEASESRESPONSE']._serialized_end=2691
+  _globals['_LEASE']._serialized_start=2694
+  _globals['_LEASE']._serialized_end=3279
+  _globals['_QUEUESERVICE']._serialized_start=3282
+  _globals['_QUEUESERVICE']._serialized_end=3698
 # @@protoc_insertion_point(module_scope)

flyte/_protos/workflow/queue_service_pb2.pyi

@@ -57,18 +57,20 @@ class TaskAction(_message.Message):
     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", "phase", "start_time", "end_time", "outputs"]
+    __slots__ = ["name", "phase", "start_time", "end_time", "outputs", "spec"]
     NAME_FIELD_NUMBER: _ClassVar[int]
     PHASE_FIELD_NUMBER: _ClassVar[int]
     START_TIME_FIELD_NUMBER: _ClassVar[int]
     END_TIME_FIELD_NUMBER: _ClassVar[int]
     OUTPUTS_FIELD_NUMBER: _ClassVar[int]
+    SPEC_FIELD_NUMBER: _ClassVar[int]
     name: str
     phase: _run_definition_pb2.Phase
     start_time: _timestamp_pb2.Timestamp
     end_time: _timestamp_pb2.Timestamp
     outputs: _run_definition_pb2.OutputReferences
-    def __init__(self, name: _Optional[str] = ..., phase: _Optional[_Union[_run_definition_pb2.Phase, str]] = ..., start_time: _Optional[_Union[_timestamp_pb2.Timestamp, _Mapping]] = ..., end_time: _Optional[_Union[_timestamp_pb2.Timestamp, _Mapping]] = ..., outputs: _Optional[_Union[_run_definition_pb2.OutputReferences, _Mapping]] = ...) -> None: ...
+    spec: _task_definition_pb2.TaskSpec
+    def __init__(self, name: _Optional[str] = ..., phase: _Optional[_Union[_run_definition_pb2.Phase, str]] = ..., start_time: _Optional[_Union[_timestamp_pb2.Timestamp, _Mapping]] = ..., end_time: _Optional[_Union[_timestamp_pb2.Timestamp, _Mapping]] = ..., outputs: _Optional[_Union[_run_definition_pb2.OutputReferences, _Mapping]] = ..., spec: _Optional[_Union[_task_definition_pb2.TaskSpec, _Mapping]] = ...) -> None: ...
 
 class ConditionAction(_message.Message):
     __slots__ = ["name", "run_id", "action_id", "type", "prompt", "description"]

flyte/_version.py

@@ -17,5 +17,5 @@ __version__: str
 __version_tuple__: VERSION_TUPLE
 version_tuple: VERSION_TUPLE
 
-__version__ = version = '2.0.0b0'
-__version_tuple__ = version_tuple = (2, 0, 0, 'b0')
+__version__ = version = '2.0.0b1'
+__version_tuple__ = version_tuple = (2, 0, 0, 'b1')

flyte/remote/_action.py

@@ -20,7 +20,6 @@ from typing import (
 import grpc
 import rich.pretty
 import rich.repr
-from google.protobuf import timestamp_pb2
 from rich.console import Console
 from rich.progress import Progress, SpinnerColumn, TextColumn, TimeElapsedColumn
 

{flyte-2.0.0b0.dist-info → flyte-2.0.0b1.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: flyte
-Version: 2.0.0b0
+Version: 2.0.0b1
 Summary: Add your description here
 Author-email: Ketan Umare <kumare3@users.noreply.github.com>
 Requires-Python: >=3.10

{flyte-2.0.0b0.dist-info → flyte-2.0.0b1.dist-info}/RECORD

@@ -25,7 +25,7 @@ flyte/_task_plugins.py,sha256=9MH3nFPOH_e8_92BT4sFk4oyAnj6GJFvaPYWaraX7yE,1037
 flyte/_timeout.py,sha256=zx5sFcbYmjJAJbZWSGzzX-BpC9HC7Jfs35T7vVhKwkk,1571
 flyte/_tools.py,sha256=tWb0sx3t3mm4jbaQVjCTc9y39oR_Ibo3z_KHToP3Lto,966
 flyte/_trace.py,sha256=SSE1nzUgmVTS2xFNtchEOjEjlRavMOIInasXzY8i9lU,4911
-flyte/_version.py,sha256=gLgPnI0TMB_aPRrdpP481xpFaTEBB7Krxu1v9wOrbJ0,519
+flyte/_version.py,sha256=tih12kl71zm-lSZ-WK20LgQ1skyymSxHAu3d0I8n6WM,519
 flyte/errors.py,sha256=3PCjGAxqifLBY6RE1n_ttmrmhcHwwtzS4UWDPql-FX4,5175
 flyte/extend.py,sha256=GB4ZedGzKa30vYWRVPOdxEeK62xnUVFY4z2tD6H9eEw,376
 flyte/models.py,sha256=k0hLoyOYQHnu5DXclEtKK2owDOXCdy0UKdGcl3xR7BQ,15250
@@ -44,11 +44,11 @@ flyte/_code_bundle/bundle.py,sha256=nUAwYTVAE3Z9dfgkBtsqCoKJImjSl4AicG36yweWHLc,
 flyte/_internal/__init__.py,sha256=vjXgGzAAjy609YFkAy9_RVPuUlslsHSJBXCLNTVnqOY,136
 flyte/_internal/controllers/__init__.py,sha256=5CBnS9lb1VFMzZuRXUiaPhlN3G9qh7Aq9kTwxW5hsRw,4301
 flyte/_internal/controllers/_local_controller.py,sha256=tbX6x8TBKwZHN1vYumftUcaX9fy162hGibo3AQZHYKA,7260
-flyte/_internal/controllers/_trace.py,sha256=NAQu0yAh2s-Ou8ZNMOXN5SxNyPmRj4h_ssxzQYklh8Q,1391
+flyte/_internal/controllers/_trace.py,sha256=ywFg_M2nGrCKYLbh4iVdsVlRtPT1K2S-XZMvDJU8AKU,1499
 flyte/_internal/controllers/remote/__init__.py,sha256=9_azH1eHLqY6VULpDugXi7Kf1kK1ODqEnsQ_3wM6IqU,1919
-flyte/_internal/controllers/remote/_action.py,sha256=YHwFIfEiTyXSDYyegVRJF5mg2joMAbB-_PIzl69S6jw,7088
+flyte/_internal/controllers/remote/_action.py,sha256=ENV1thRXllSpi2s4idL-vCVcmXQNS17hmP2lMDKzNdo,7397
 flyte/_internal/controllers/remote/_client.py,sha256=HPbzbfaWZVv5wpOvKNtFXR6COiZDwd1cUJQqi60A7oU,1421
-flyte/_internal/controllers/remote/_controller.py,sha256=EDzALJnLQWRMxGnFolVrI7xhBszo5apv2ftal2UuCgA,23520
+flyte/_internal/controllers/remote/_controller.py,sha256=oubfymmNgZtVEDJF8HiV3r2i_6lGgBZBzh0437jY0zQ,23766
 flyte/_internal/controllers/remote/_core.py,sha256=49l1X3YFyX-QYeSlvkOWdyEU2xc1Fr7I8qHvC3nQ6MA,19069
 flyte/_internal/controllers/remote/_informer.py,sha256=w4p29_dzS_ns762eNBljvnbJLgCm36d1Ogo2ZkgV1yg,14418
 flyte/_internal/controllers/remote/_service_protocol.py,sha256=B9qbIg6DiGeac-iSccLmX_AL2xUgX4ezNUOiAbSy4V0,1357
@@ -124,8 +124,8 @@ flyte/_protos/workflow/environment_pb2_grpc.py,sha256=1oboBPFxaTEXt9Aw7EAj8gXHDC
 flyte/_protos/workflow/node_execution_service_pb2.py,sha256=IOLg3tNikY7n00kLOVsC69yyXc5Ttnx-_-xUuc0q05Q,1654
 flyte/_protos/workflow/node_execution_service_pb2.pyi,sha256=C7VVuw_bnxp68qemD3SLoGIL-Hmno6qkIoq3l6W2qb8,135
 flyte/_protos/workflow/node_execution_service_pb2_grpc.py,sha256=2JJDS3Aww3FFDW-qYdTaxC75gRpsgnn4an6LPZmF9uA,947
-flyte/_protos/workflow/queue_service_pb2.py,sha256=0KL5mQNeA2-w4ceSZDj7vP-hoHKoSCYhkykhDCPoU4U,13297
-flyte/_protos/workflow/queue_service_pb2.pyi,sha256=GiVvoiRgVtvf_l3_tfIEg4DesyVFnq4r7UBm_g5KOYg,9341
+flyte/_protos/workflow/queue_service_pb2.py,sha256=zrYxvapY74-xz0RVXsXv6wG0u054kUnNQDL5b7GJ5lw,13561
+flyte/_protos/workflow/queue_service_pb2.pyi,sha256=lZ9rO0KpuyGadl5GqQ8JgffZcDlzfSeyNmsZ3du9Gqg,9499
 flyte/_protos/workflow/queue_service_pb2_grpc.py,sha256=6KK87jYXrmK0jacf4AKhHp21QU9JFJPOiEBjbDRkBm0,7839
 flyte/_protos/workflow/run_definition_pb2.py,sha256=sdyl7xgl6tTqkKNXrWcrhvV4KOZPqoMa9vjEmwVyfZc,16193
 flyte/_protos/workflow/run_definition_pb2.pyi,sha256=xQ5bXt3J0hoXUGTqEvAGqm0fJT7dYOAl1DSbX83MrD4,16930
@@ -184,7 +184,7 @@ flyte/migrate/dynamic.py,sha256=igL6lMgKKH-g3CwF48SGi-F5YWdE9p8iXlMi4LJGnxc,296
 flyte/migrate/task.py,sha256=V1-bKQuz0aPYG9jzHpjdujPJSGTQRH4Yda3ZPmwkVxA,3714
 flyte/migrate/workflow.py,sha256=AbPchkUZXUAAGlrdy43M23fdpo9tsVaXcMLVSh7z9pU,328
 flyte/remote/__init__.py,sha256=y9eke9JzEJkygk8eKZjSj44fJGlyepuW4i-j6lbCAPY,617
-flyte/remote/_action.py,sha256=_EGMSjuqYAkxSrproa2P7107WuHe_ESN2lnJlkBBdhI,23881
+flyte/remote/_action.py,sha256=5DN6rEb4WhsdDEZsYPFFOUyn2yjSMefVojkMF7DE7XM,23839
 flyte/remote/_console.py,sha256=avmELJPx8nQMAVPrHlh6jEIRPjrMwFpdZjJsWOOa9rE,660
 flyte/remote/_data.py,sha256=zYXXlqEvPdsC44Gm7rP_hQjRgVe3EFfhZNEWKF0p4MQ,6163
 flyte/remote/_logs.py,sha256=aDG18-uPVb2J3PxmqmAY1C0Z4Iv1P1agg-iF4nSQR4U,6709
@@ -230,10 +230,10 @@ flyte/types/_renderer.py,sha256=ygcCo5l60lHufyQISFddZfWwLlQ8kJAKxUT_XnR_6dY,4818
 flyte/types/_string_literals.py,sha256=NlG1xV8RSA-sZ-n-IFQCAsdB6jXJOAKkHWtnopxVVDk,4231
 flyte/types/_type_engine.py,sha256=Tas_OXYddOi0nDuORjqan2SkJ96wKD8937I2l1bo8vk,97916
 flyte/types/_utils.py,sha256=pbts9E1_2LTdLygAY0UYTLYJ8AsN3BZyviSXvrtcutc,2626
-flyte-2.0.0b0.data/scripts/runtime.py,sha256=2jTy3ccvrJ__Xrfdo2t0Fxhsojc5o2zIxDHt98RE_eU,6475
-flyte-2.0.0b0.dist-info/licenses/LICENSE,sha256=xx0jnfkXJvxRnG63LTGOxlggYnIysveWIZ6H3PNdCrQ,11357
-flyte-2.0.0b0.dist-info/METADATA,sha256=dJhOX_qOoMQVGvkCFJDl8c26Jr60yVgcrwVFYiloEiM,10015
-flyte-2.0.0b0.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
-flyte-2.0.0b0.dist-info/entry_points.txt,sha256=MIq2z5dBurdCJfpXfMKzgBv7sJOakKRYxr8G0cMiTrg,75
-flyte-2.0.0b0.dist-info/top_level.txt,sha256=7dkyFbikvA12LEZEqawx8oDG1CMod6hTliPj7iWzgYo,6
-flyte-2.0.0b0.dist-info/RECORD,,
+flyte-2.0.0b1.data/scripts/runtime.py,sha256=2jTy3ccvrJ__Xrfdo2t0Fxhsojc5o2zIxDHt98RE_eU,6475
+flyte-2.0.0b1.dist-info/licenses/LICENSE,sha256=xx0jnfkXJvxRnG63LTGOxlggYnIysveWIZ6H3PNdCrQ,11357
+flyte-2.0.0b1.dist-info/METADATA,sha256=v8ZPND_lTS1evo9ypHWHFL7pcb1Xs1xB3b5eW_d9WNs,10015
+flyte-2.0.0b1.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
+flyte-2.0.0b1.dist-info/entry_points.txt,sha256=MIq2z5dBurdCJfpXfMKzgBv7sJOakKRYxr8G0cMiTrg,75
+flyte-2.0.0b1.dist-info/top_level.txt,sha256=7dkyFbikvA12LEZEqawx8oDG1CMod6hTliPj7iWzgYo,6
+flyte-2.0.0b1.dist-info/RECORD,,