flyte 0.2.0b3__py3-none-any.whl → 0.2.0b5__py3-none-any.whl

flyte/_build.py

@@ -1,10 +1,11 @@
 from __future__ import annotations
 
-from ._api_commons import syncer
+from flyte.syncify import syncify
+
 from ._image import Image
 
 
-@syncer.wrap
+@syncify
 async def build(image: Image) -> str:
     """
     Build an image. The existing async context will be used.

flyte/_deploy.py

@@ -7,11 +7,11 @@ from typing import TYPE_CHECKING, Dict, List, Optional, Tuple
 import rich.repr
 
 from flyte.models import SerializationContext
+from flyte.syncify import syncify
 
-from ._api_commons import syncer
 from ._environment import Environment
 from ._image import Image
-from ._initialize import get_client, get_common_config, requires_client, requires_initialization
+from ._initialize import ensure_client, get_client, get_common_config, requires_initialization
 from ._logging import logger
 from ._task import TaskTemplate
 from ._task_environment import TaskEnvironment
@@ -47,13 +47,13 @@ class Deployment:
         return f"Deployment(envs=[{env_names}], tasks=[{task_names_versions}])"
 
 
-@requires_client
 async def _deploy_task(
     task: TaskTemplate, serialization_context: SerializationContext, dryrun: bool = False
 ) -> task_definition_pb2.TaskSpec:
     """
     Deploy the given task.
     """
+    ensure_client()
     from ._internal.runtime.task_serde import translate_task_to_wire
     from ._protos.workflow import task_definition_pb2, task_service_pb2
 
@@ -177,7 +177,7 @@ def plan_deploy(*envs: Environment, version: Optional[str] = None) -> Deployment
     return DeploymentPlan(planned_envs, version=version)
 
 
-@syncer.wrap
+@syncify
 async def deploy(
     *envs: Environment,
     dryrun: bool = False,

flyte/_initialize.py

@@ -8,8 +8,8 @@ from pathlib import Path
 from typing import TYPE_CHECKING, Callable, List, Literal, Optional, TypeVar
 
 from flyte.errors import InitializationError
+from flyte.syncify import syncify
 
-from ._api_commons import syncer
 from ._logging import initialize_logger
 from ._tools import ipython_check
 
@@ -108,7 +108,7 @@ async def _initialize_client(
     )
 
 
-@syncer.wrap
+@syncify
 async def init(
     org: str | None = None,
     project: str | None = None,
@@ -298,6 +298,20 @@ def initialize_in_cluster() -> None:
 T = TypeVar("T", bound=Callable)
 
 
+def ensure_client():
+    """
+    Ensure that the client is initialized. If not, raise an InitializationError.
+    This function is used to check if the client is initialized before executing any Flyte remote API methods.
+    """
+    if _get_init_config() is None or _get_init_config().client is None:
+        raise InitializationError(
+            "ClientNotInitializedError",
+            "user",
+            "Client has not been initialized. Call flyte.init() with a valid endpoint"
+            " or api-key before using this function.",
+        )
+
+
 def requires_client(func: T) -> T:
     """
     Decorator that checks if the client has been initialized before executing the function.
@@ -308,7 +322,7 @@ def requires_client(func: T) -> T:
     """
 
     @functools.wraps(func)
-    def wrapper(*args, **kwargs) -> T:
+    async def wrapper(*args, **kwargs) -> T:
         init_config = _get_init_config()
         if init_config is None or init_config.client is None:
             raise InitializationError(

flyte/_internal/controllers/remote/_core.py

@@ -286,10 +286,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]:

flyte/_internal/controllers/remote/_service_protocol.py

@@ -28,12 +28,12 @@ class QueueService(Protocol):
     ) -> queue_service_pb2.EnqueueActionResponse:
         """Enqueue a task"""
 
-    async def AbortQueuedAction(
-        self,
-        req: queue_service_pb2.AbortQueuedActionRequest,
-        **kwargs,
-    ) -> queue_service_pb2.AbortQueuedActionResponse:
-        """Dequeue a task"""
+    # async def AbortQueuedAction(
+    #     self,
+    #     req: queue_service_pb2.AbortQueuedActionRequest,
+    #     **kwargs,
+    # ) -> queue_service_pb2.AbortQueuedActionResponse:
+    #     """Dequeue a task"""
 
 
 class ClientSet(Protocol):

flyte/_protos/logs/dataplane/payload_pb2.py

@@ -15,7 +15,7 @@ from google.protobuf import timestamp_pb2 as google_dot_protobuf_dot_timestamp__
 from flyte._protos.validate.validate import validate_pb2 as validate_dot_validate__pb2
 
 
-DESCRIPTOR = _descriptor_pool.Default().AddSerializedFile(b'\n\x1clogs/dataplane/payload.proto\x12\x17\x63loudidl.logs.dataplane\x1a\x1fgoogle/protobuf/timestamp.proto\x1a\x17validate/validate.proto\"o\n\x0bPodResource\x12%\n\tnamespace\x18\x01 \x01(\tB\x07\xfa\x42\x04r\x02\x10\x01R\tnamespace\x12\x1b\n\x04name\x18\x02 \x01(\tB\x07\xfa\x42\x04r\x02\x10\x01R\x04name\x12\x1c\n\tcontainer\x18\x03 \x01(\tR\tcontainer\"\xda\x01\n\x1cTailTaskExecutionLogsRequest\x12Z\n\x0flogging_context\x18\x01 \x01(\x0b\x32\'.cloudidl.logs.dataplane.LoggingContextB\x08\xfa\x42\x05\x8a\x01\x02\x10\x01R\x0eloggingContext\x12;\n\x06source\x18\x03 \x01(\x0e\x32#.cloudidl.logs.dataplane.LogsSourceR\x06source\x12\x1b\n\tno_follow\x18\x04 \x01(\x08R\x08noFollowJ\x04\x08\x02\x10\x03\"\xeb\x01\n\x1dTailTaskExecutionLogsResponse\x12Q\n\x04logs\x18\x01 \x01(\x0b\x32;.cloudidl.logs.dataplane.TailTaskExecutionLogsResponse.LogsH\x00R\x04logs\x1am\n\x04Logs\x12\x18\n\x05lines\x18\x01 \x03(\tB\x02\x18\x01R\x05lines\x12K\n\x10structured_lines\x18\x02 \x03(\x0b\x32 .cloudidl.logs.dataplane.LogLineR\x0fstructuredLinesB\x08\n\x06result\"\xb8\x03\n\x0eLoggingContext\x12*\n\x0c\x63luster_name\x18\x03 \x01(\tB\x07\xfa\x42\x04r\x02\x10\x01R\x0b\x63lusterName\x12:\n\x14kubernetes_namespace\x18\x04 \x01(\tB\x07\xfa\x42\x04r\x02\x10\x01R\x13kubernetesNamespace\x12\x37\n\x13kubernetes_pod_name\x18\x05 \x01(\tB\x07\xfa\x42\x04r\x02\x10\x01R\x11kubernetesPodName\x12\x43\n\x19kubernetes_container_name\x18\x06 \x01(\tB\x07\xfa\x42\x04r\x02\x10\x01R\x17kubernetesContainerName\x12[\n\x1c\x65xecution_attempt_start_time\x18\x07 \x01(\x0b\x32\x1a.google.protobuf.TimestampR\x19\x65xecutionAttemptStartTime\x12W\n\x1a\x65xecution_attempt_end_time\x18\x08 \x01(\x0b\x32\x1a.google.protobuf.TimestampR\x17\x65xecutionAttemptEndTimeJ\x04\x08\x01\x10\x02J\x04\x08\x02\x10\x03\"\xf2\x01\n\x13\x43ontainerIdentifier\x12*\n\x0c\x63luster_name\x18\x01 \x01(\tB\x07\xfa\x42\x04r\x02\x10\x01R\x0b\x63lusterName\x12:\n\x14kubernetes_namespace\x18\x02 \x01(\tB\x07\xfa\x42\x04r\x02\x10\x01R\x13kubernetesNamespace\x12\x37\n\x13kubernetes_pod_name\x18\x03 \x01(\tB\x07\xfa\x42\x04r\x02\x10\x01R\x11kubernetesPodName\x12:\n\x19kubernetes_container_name\x18\x04 \x01(\tR\x17kubernetesContainerName\"\xb7\x02\n\x11\x43ontainerSelector\x12*\n\x0c\x63luster_name\x18\x01 \x01(\tB\x07\xfa\x42\x04r\x02\x10\x01R\x0b\x63lusterName\x12:\n\x14kubernetes_namespace\x18\x02 \x01(\tB\x07\xfa\x42\x04r\x02\x10\x01R\x13kubernetesNamespace\x12;\n\x1akubernetes_pod_name_prefix\x18\x03 \x01(\tR\x17kubernetesPodNamePrefix\x12:\n\x19kubernetes_container_name\x18\x04 \x01(\tR\x17kubernetesContainerName\x12\x41\n\x1dkubernetes_pod_label_selector\x18\x05 \x01(\tR\x1akubernetesPodLabelSelector\"^\n\x0fLiveLogsOptions\x12$\n\x0elog_pod_status\x18\x01 \x01(\x08R\x0clogPodStatus\x12%\n\x0elog_timestamps\x18\x02 \x01(\x08R\rlogTimestamps\"\xd5\x03\n\x0fTailLogsRequest\x12L\n\tcontainer\x18\x01 \x01(\x0b\x32,.cloudidl.logs.dataplane.ContainerIdentifierH\x00R\tcontainer\x12[\n\x12\x63ontainer_selector\x18\x02 \x01(\x0b\x32*.cloudidl.logs.dataplane.ContainerSelectorH\x00R\x11\x63ontainerSelector\x12\x43\n\nstart_time\x18\x03 \x01(\x0b\x32\x1a.google.protobuf.TimestampB\x08\xfa\x42\x05\xb2\x01\x02\x08\x01R\tstartTime\x12\x35\n\x08\x65nd_time\x18\x04 \x01(\x0b\x32\x1a.google.protobuf.TimestampR\x07\x65ndTime\x12;\n\x06source\x18\x05 \x01(\x0e\x32#.cloudidl.logs.dataplane.LogsSourceR\x06source\x12T\n\x11live_logs_options\x18\x06 \x01(\x0b\x32(.cloudidl.logs.dataplane.LiveLogsOptionsR\x0fliveLogsOptionsB\x08\n\x06target\"\xa9\x01\n\x07LogLine\x12\x38\n\ttimestamp\x18\x01 \x01(\x0b\x32\x1a.google.protobuf.TimestampR\ttimestamp\x12\x18\n\x07message\x18\x02 \x01(\tR\x07message\x12J\n\noriginator\x18\x03 \x01(\x0e\x32*.cloudidl.logs.dataplane.LogLineOriginatorR\noriginator\"\xf0\x01\n\x08LogLines\x12\x18\n\x05lines\x18\x01 \x03(\tB\x02\x18\x01R\x05lines\x12\'\n\x0f\x63ontainer_index\x18\x02 \x01(\rR\x0e\x63ontainerIndex\x12T\n\tcontainer\x18\x03 \x01(\x0b\x32,.cloudidl.logs.dataplane.ContainerIdentifierB\x08\xfa\x42\x05\x8a\x01\x02\x10\x01R\tcontainer\x12K\n\x10structured_lines\x18\x04 \x03(\x0b\x32 .cloudidl.logs.dataplane.LogLineR\x0fstructuredLines\"a\n\x11LogContainersList\x12L\n\ncontainers\x18\x01 \x03(\x0b\x32,.cloudidl.logs.dataplane.ContainerIdentifierR\ncontainers\"F\n\rLogLinesBatch\x12\x35\n\x04logs\x18\x01 \x03(\x0b\x32!.cloudidl.logs.dataplane.LogLinesR\x04logs\"\xd9\x01\n\x10TailLogsResponse\x12V\n\ncontainers\x18\x01 \x01(\x0b\x32*.cloudidl.logs.dataplane.LogContainersListB\x08\xfa\x42\x05\x8a\x01\x02\x10\x01H\x00R\ncontainers\x12Z\n\x0flog_lines_batch\x18\x03 \x01(\x0b\x32&.cloudidl.logs.dataplane.LogLinesBatchB\x08\xfa\x42\x05\x8a\x01\x02\x10\x01H\x00R\rlogLinesBatchB\x0b\n\x04resp\x12\x03\xf8\x42\x01J\x04\x08\x02\x10\x03*6\n\x11LogLineOriginator\x12\x0b\n\x07UNKNOWN\x10\x00\x12\x08\n\x04USER\x10\x01\x12\n\n\x06SYSTEM\x10\x02*F\n\nLogsSource\x12\x15\n\x11LIVE_OR_PERSISTED\x10\x00\x12\r\n\tLIVE_ONLY\x10\x01\x12\x12\n\x0ePERSISTED_ONLY\x10\x02\x42\xde\x01\n\x1b\x63om.cloudidl.logs.dataplaneB\x0cPayloadProtoH\x02P\x01Z1github.com/unionai/cloud/gen/pb-go/logs/dataplane\xa2\x02\x03\x43LD\xaa\x02\x17\x43loudidl.Logs.Dataplane\xca\x02\x17\x43loudidl\\Logs\\Dataplane\xe2\x02#Cloudidl\\Logs\\Dataplane\\GPBMetadata\xea\x02\x19\x43loudidl::Logs::Dataplaneb\x06proto3')
+DESCRIPTOR = _descriptor_pool.Default().AddSerializedFile(b'\n\x1clogs/dataplane/payload.proto\x12\x17\x63loudidl.logs.dataplane\x1a\x1fgoogle/protobuf/timestamp.proto\x1a\x17validate/validate.proto\"o\n\x0bPodResource\x12%\n\tnamespace\x18\x01 \x01(\tB\x07\xfa\x42\x04r\x02\x10\x01R\tnamespace\x12\x1b\n\x04name\x18\x02 \x01(\tB\x07\xfa\x42\x04r\x02\x10\x01R\x04name\x12\x1c\n\tcontainer\x18\x03 \x01(\tR\tcontainer\"\xda\x01\n\x1cTailTaskExecutionLogsRequest\x12Z\n\x0flogging_context\x18\x01 \x01(\x0b\x32\'.cloudidl.logs.dataplane.LoggingContextB\x08\xfa\x42\x05\x8a\x01\x02\x10\x01R\x0eloggingContext\x12;\n\x06source\x18\x03 \x01(\x0e\x32#.cloudidl.logs.dataplane.LogsSourceR\x06source\x12\x1b\n\tno_follow\x18\x04 \x01(\x08R\x08noFollowJ\x04\x08\x02\x10\x03\"\xeb\x01\n\x1dTailTaskExecutionLogsResponse\x12Q\n\x04logs\x18\x01 \x01(\x0b\x32;.cloudidl.logs.dataplane.TailTaskExecutionLogsResponse.LogsH\x00R\x04logs\x1am\n\x04Logs\x12\x18\n\x05lines\x18\x01 \x03(\tB\x02\x18\x01R\x05lines\x12K\n\x10structured_lines\x18\x02 \x03(\x0b\x32 .cloudidl.logs.dataplane.LogLineR\x0fstructuredLinesB\x08\n\x06result\"\xf6\x04\n\x0eLoggingContext\x12*\n\x0c\x63luster_name\x18\x03 \x01(\tB\x07\xfa\x42\x04r\x02\x10\x01R\x0b\x63lusterName\x12:\n\x14kubernetes_namespace\x18\x04 \x01(\tB\x07\xfa\x42\x04r\x02\x10\x01R\x13kubernetesNamespace\x12\x37\n\x13kubernetes_pod_name\x18\x05 \x01(\tB\x07\xfa\x42\x04r\x02\x10\x01R\x11kubernetesPodName\x12\x43\n\x19kubernetes_container_name\x18\x06 \x01(\tB\x07\xfa\x42\x04r\x02\x10\x01R\x17kubernetesContainerName\x12[\n\x1c\x65xecution_attempt_start_time\x18\x07 \x01(\x0b\x32\x1a.google.protobuf.TimestampR\x19\x65xecutionAttemptStartTime\x12W\n\x1a\x65xecution_attempt_end_time\x18\x08 \x01(\x0b\x32\x1a.google.protobuf.TimestampR\x17\x65xecutionAttemptEndTime\x12t\n\x15kubernetes_pod_labels\x18\t \x03(\x0b\x32@.cloudidl.logs.dataplane.LoggingContext.KubernetesPodLabelsEntryR\x13kubernetesPodLabels\x1a\x46\n\x18KubernetesPodLabelsEntry\x12\x10\n\x03key\x18\x01 \x01(\tR\x03key\x12\x14\n\x05value\x18\x02 \x01(\tR\x05value:\x02\x38\x01J\x04\x08\x01\x10\x02J\x04\x08\x02\x10\x03\"\xf2\x01\n\x13\x43ontainerIdentifier\x12*\n\x0c\x63luster_name\x18\x01 \x01(\tB\x07\xfa\x42\x04r\x02\x10\x01R\x0b\x63lusterName\x12:\n\x14kubernetes_namespace\x18\x02 \x01(\tB\x07\xfa\x42\x04r\x02\x10\x01R\x13kubernetesNamespace\x12\x37\n\x13kubernetes_pod_name\x18\x03 \x01(\tB\x07\xfa\x42\x04r\x02\x10\x01R\x11kubernetesPodName\x12:\n\x19kubernetes_container_name\x18\x04 \x01(\tR\x17kubernetesContainerName\"\xb7\x02\n\x11\x43ontainerSelector\x12*\n\x0c\x63luster_name\x18\x01 \x01(\tB\x07\xfa\x42\x04r\x02\x10\x01R\x0b\x63lusterName\x12:\n\x14kubernetes_namespace\x18\x02 \x01(\tB\x07\xfa\x42\x04r\x02\x10\x01R\x13kubernetesNamespace\x12;\n\x1akubernetes_pod_name_prefix\x18\x03 \x01(\tR\x17kubernetesPodNamePrefix\x12:\n\x19kubernetes_container_name\x18\x04 \x01(\tR\x17kubernetesContainerName\x12\x41\n\x1dkubernetes_pod_label_selector\x18\x05 \x01(\tR\x1akubernetesPodLabelSelector\"^\n\x0fLiveLogsOptions\x12$\n\x0elog_pod_status\x18\x01 \x01(\x08R\x0clogPodStatus\x12%\n\x0elog_timestamps\x18\x02 \x01(\x08R\rlogTimestamps\"\xd5\x03\n\x0fTailLogsRequest\x12L\n\tcontainer\x18\x01 \x01(\x0b\x32,.cloudidl.logs.dataplane.ContainerIdentifierH\x00R\tcontainer\x12[\n\x12\x63ontainer_selector\x18\x02 \x01(\x0b\x32*.cloudidl.logs.dataplane.ContainerSelectorH\x00R\x11\x63ontainerSelector\x12\x43\n\nstart_time\x18\x03 \x01(\x0b\x32\x1a.google.protobuf.TimestampB\x08\xfa\x42\x05\xb2\x01\x02\x08\x01R\tstartTime\x12\x35\n\x08\x65nd_time\x18\x04 \x01(\x0b\x32\x1a.google.protobuf.TimestampR\x07\x65ndTime\x12;\n\x06source\x18\x05 \x01(\x0e\x32#.cloudidl.logs.dataplane.LogsSourceR\x06source\x12T\n\x11live_logs_options\x18\x06 \x01(\x0b\x32(.cloudidl.logs.dataplane.LiveLogsOptionsR\x0fliveLogsOptionsB\x08\n\x06target\"\xa9\x01\n\x07LogLine\x12\x38\n\ttimestamp\x18\x01 \x01(\x0b\x32\x1a.google.protobuf.TimestampR\ttimestamp\x12\x18\n\x07message\x18\x02 \x01(\tR\x07message\x12J\n\noriginator\x18\x03 \x01(\x0e\x32*.cloudidl.logs.dataplane.LogLineOriginatorR\noriginator\"\xf0\x01\n\x08LogLines\x12\x18\n\x05lines\x18\x01 \x03(\tB\x02\x18\x01R\x05lines\x12\'\n\x0f\x63ontainer_index\x18\x02 \x01(\rR\x0e\x63ontainerIndex\x12T\n\tcontainer\x18\x03 \x01(\x0b\x32,.cloudidl.logs.dataplane.ContainerIdentifierB\x08\xfa\x42\x05\x8a\x01\x02\x10\x01R\tcontainer\x12K\n\x10structured_lines\x18\x04 \x03(\x0b\x32 .cloudidl.logs.dataplane.LogLineR\x0fstructuredLines\"a\n\x11LogContainersList\x12L\n\ncontainers\x18\x01 \x03(\x0b\x32,.cloudidl.logs.dataplane.ContainerIdentifierR\ncontainers\"F\n\rLogLinesBatch\x12\x35\n\x04logs\x18\x01 \x03(\x0b\x32!.cloudidl.logs.dataplane.LogLinesR\x04logs\"\xd9\x01\n\x10TailLogsResponse\x12V\n\ncontainers\x18\x01 \x01(\x0b\x32*.cloudidl.logs.dataplane.LogContainersListB\x08\xfa\x42\x05\x8a\x01\x02\x10\x01H\x00R\ncontainers\x12Z\n\x0flog_lines_batch\x18\x03 \x01(\x0b\x32&.cloudidl.logs.dataplane.LogLinesBatchB\x08\xfa\x42\x05\x8a\x01\x02\x10\x01H\x00R\rlogLinesBatchB\x0b\n\x04resp\x12\x03\xf8\x42\x01J\x04\x08\x02\x10\x03*6\n\x11LogLineOriginator\x12\x0b\n\x07UNKNOWN\x10\x00\x12\x08\n\x04USER\x10\x01\x12\n\n\x06SYSTEM\x10\x02*F\n\nLogsSource\x12\x15\n\x11LIVE_OR_PERSISTED\x10\x00\x12\r\n\tLIVE_ONLY\x10\x01\x12\x12\n\x0ePERSISTED_ONLY\x10\x02\x42\xde\x01\n\x1b\x63om.cloudidl.logs.dataplaneB\x0cPayloadProtoH\x02P\x01Z1github.com/unionai/cloud/gen/pb-go/logs/dataplane\xa2\x02\x03\x43LD\xaa\x02\x17\x43loudidl.Logs.Dataplane\xca\x02\x17\x43loudidl\\Logs\\Dataplane\xe2\x02#Cloudidl\\Logs\\Dataplane\\GPBMetadata\xea\x02\x19\x43loudidl::Logs::Dataplaneb\x06proto3')
 
 _globals = globals()
 _builder.BuildMessageAndEnumDescriptors(DESCRIPTOR, _globals)
@@ -31,6 +31,8 @@ if _descriptor._USE_C_DESCRIPTORS == False:
   _TAILTASKEXECUTIONLOGSREQUEST.fields_by_name['logging_context']._serialized_options = b'\372B\005\212\001\002\020\001'
   _TAILTASKEXECUTIONLOGSRESPONSE_LOGS.fields_by_name['lines']._options = None
   _TAILTASKEXECUTIONLOGSRESPONSE_LOGS.fields_by_name['lines']._serialized_options = b'\030\001'
+  _LOGGINGCONTEXT_KUBERNETESPODLABELSENTRY._options = None
+  _LOGGINGCONTEXT_KUBERNETESPODLABELSENTRY._serialized_options = b'8\001'
   _LOGGINGCONTEXT.fields_by_name['cluster_name']._options = None
   _LOGGINGCONTEXT.fields_by_name['cluster_name']._serialized_options = b'\372B\004r\002\020\001'
   _LOGGINGCONTEXT.fields_by_name['kubernetes_namespace']._options = None
@@ -61,10 +63,10 @@ if _descriptor._USE_C_DESCRIPTORS == False:
   _TAILLOGSRESPONSE.fields_by_name['containers']._serialized_options = b'\372B\005\212\001\002\020\001'
   _TAILLOGSRESPONSE.fields_by_name['log_lines_batch']._options = None
   _TAILLOGSRESPONSE.fields_by_name['log_lines_batch']._serialized_options = b'\372B\005\212\001\002\020\001'
-  _globals['_LOGLINEORIGINATOR']._serialized_start=3063
-  _globals['_LOGLINEORIGINATOR']._serialized_end=3117
-  _globals['_LOGSSOURCE']._serialized_start=3119
-  _globals['_LOGSSOURCE']._serialized_end=3189
+  _globals['_LOGLINEORIGINATOR']._serialized_start=3253
+  _globals['_LOGLINEORIGINATOR']._serialized_end=3307
+  _globals['_LOGSSOURCE']._serialized_start=3309
+  _globals['_LOGSSOURCE']._serialized_end=3379
   _globals['_PODRESOURCE']._serialized_start=115
   _globals['_PODRESOURCE']._serialized_end=226
   _globals['_TAILTASKEXECUTIONLOGSREQUEST']._serialized_start=229
@@ -74,23 +76,25 @@ if _descriptor._USE_C_DESCRIPTORS == False:
   _globals['_TAILTASKEXECUTIONLOGSRESPONSE_LOGS']._serialized_start=566
   _globals['_TAILTASKEXECUTIONLOGSRESPONSE_LOGS']._serialized_end=675
   _globals['_LOGGINGCONTEXT']._serialized_start=688
-  _globals['_LOGGINGCONTEXT']._serialized_end=1128
-  _globals['_CONTAINERIDENTIFIER']._serialized_start=1131
-  _globals['_CONTAINERIDENTIFIER']._serialized_end=1373
-  _globals['_CONTAINERSELECTOR']._serialized_start=1376
-  _globals['_CONTAINERSELECTOR']._serialized_end=1687
-  _globals['_LIVELOGSOPTIONS']._serialized_start=1689
-  _globals['_LIVELOGSOPTIONS']._serialized_end=1783
-  _globals['_TAILLOGSREQUEST']._serialized_start=1786
-  _globals['_TAILLOGSREQUEST']._serialized_end=2255
-  _globals['_LOGLINE']._serialized_start=2258
-  _globals['_LOGLINE']._serialized_end=2427
-  _globals['_LOGLINES']._serialized_start=2430
-  _globals['_LOGLINES']._serialized_end=2670
-  _globals['_LOGCONTAINERSLIST']._serialized_start=2672
-  _globals['_LOGCONTAINERSLIST']._serialized_end=2769
-  _globals['_LOGLINESBATCH']._serialized_start=2771
-  _globals['_LOGLINESBATCH']._serialized_end=2841
-  _globals['_TAILLOGSRESPONSE']._serialized_start=2844
-  _globals['_TAILLOGSRESPONSE']._serialized_end=3061
+  _globals['_LOGGINGCONTEXT']._serialized_end=1318
+  _globals['_LOGGINGCONTEXT_KUBERNETESPODLABELSENTRY']._serialized_start=1236
+  _globals['_LOGGINGCONTEXT_KUBERNETESPODLABELSENTRY']._serialized_end=1306
+  _globals['_CONTAINERIDENTIFIER']._serialized_start=1321
+  _globals['_CONTAINERIDENTIFIER']._serialized_end=1563
+  _globals['_CONTAINERSELECTOR']._serialized_start=1566
+  _globals['_CONTAINERSELECTOR']._serialized_end=1877
+  _globals['_LIVELOGSOPTIONS']._serialized_start=1879
+  _globals['_LIVELOGSOPTIONS']._serialized_end=1973
+  _globals['_TAILLOGSREQUEST']._serialized_start=1976
+  _globals['_TAILLOGSREQUEST']._serialized_end=2445
+  _globals['_LOGLINE']._serialized_start=2448
+  _globals['_LOGLINE']._serialized_end=2617
+  _globals['_LOGLINES']._serialized_start=2620
+  _globals['_LOGLINES']._serialized_end=2860
+  _globals['_LOGCONTAINERSLIST']._serialized_start=2862
+  _globals['_LOGCONTAINERSLIST']._serialized_end=2959
+  _globals['_LOGLINESBATCH']._serialized_start=2961
+  _globals['_LOGLINESBATCH']._serialized_end=3031
+  _globals['_TAILLOGSRESPONSE']._serialized_start=3034
+  _globals['_TAILLOGSRESPONSE']._serialized_end=3251
 # @@protoc_insertion_point(module_scope)

flyte/_protos/logs/dataplane/payload_pb2.pyi

@@ -60,20 +60,29 @@ class TailTaskExecutionLogsResponse(_message.Message):
     def __init__(self, logs: _Optional[_Union[TailTaskExecutionLogsResponse.Logs, _Mapping]] = ...) -> None: ...
 
 class LoggingContext(_message.Message):
-    __slots__ = ["cluster_name", "kubernetes_namespace", "kubernetes_pod_name", "kubernetes_container_name", "execution_attempt_start_time", "execution_attempt_end_time"]
+    __slots__ = ["cluster_name", "kubernetes_namespace", "kubernetes_pod_name", "kubernetes_container_name", "execution_attempt_start_time", "execution_attempt_end_time", "kubernetes_pod_labels"]
+    class KubernetesPodLabelsEntry(_message.Message):
+        __slots__ = ["key", "value"]
+        KEY_FIELD_NUMBER: _ClassVar[int]
+        VALUE_FIELD_NUMBER: _ClassVar[int]
+        key: str
+        value: str
+        def __init__(self, key: _Optional[str] = ..., value: _Optional[str] = ...) -> None: ...
     CLUSTER_NAME_FIELD_NUMBER: _ClassVar[int]
     KUBERNETES_NAMESPACE_FIELD_NUMBER: _ClassVar[int]
     KUBERNETES_POD_NAME_FIELD_NUMBER: _ClassVar[int]
     KUBERNETES_CONTAINER_NAME_FIELD_NUMBER: _ClassVar[int]
     EXECUTION_ATTEMPT_START_TIME_FIELD_NUMBER: _ClassVar[int]
     EXECUTION_ATTEMPT_END_TIME_FIELD_NUMBER: _ClassVar[int]
+    KUBERNETES_POD_LABELS_FIELD_NUMBER: _ClassVar[int]
     cluster_name: str
     kubernetes_namespace: str
     kubernetes_pod_name: str
     kubernetes_container_name: str
     execution_attempt_start_time: _timestamp_pb2.Timestamp
     execution_attempt_end_time: _timestamp_pb2.Timestamp
-    def __init__(self, cluster_name: _Optional[str] = ..., kubernetes_namespace: _Optional[str] = ..., kubernetes_pod_name: _Optional[str] = ..., kubernetes_container_name: _Optional[str] = ..., execution_attempt_start_time: _Optional[_Union[_timestamp_pb2.Timestamp, _Mapping]] = ..., execution_attempt_end_time: _Optional[_Union[_timestamp_pb2.Timestamp, _Mapping]] = ...) -> None: ...
+    kubernetes_pod_labels: _containers.ScalarMap[str, str]
+    def __init__(self, cluster_name: _Optional[str] = ..., kubernetes_namespace: _Optional[str] = ..., kubernetes_pod_name: _Optional[str] = ..., kubernetes_container_name: _Optional[str] = ..., execution_attempt_start_time: _Optional[_Union[_timestamp_pb2.Timestamp, _Mapping]] = ..., execution_attempt_end_time: _Optional[_Union[_timestamp_pb2.Timestamp, _Mapping]] = ..., kubernetes_pod_labels: _Optional[_Mapping[str, str]] = ...) -> None: ...
 
 class ContainerIdentifier(_message.Message):
     __slots__ = ["cluster_name", "kubernetes_namespace", "kubernetes_pod_name", "kubernetes_container_name"]

flyte/_protos/workflow/common_pb2.py

@@ -0,0 +1,27 @@
+# -*- coding: utf-8 -*-
+# Generated by the protocol buffer compiler.  DO NOT EDIT!
+# source: workflow/common.proto
+"""Generated protocol buffer code."""
+from google.protobuf import descriptor as _descriptor
+from google.protobuf import descriptor_pool as _descriptor_pool
+from google.protobuf import symbol_database as _symbol_database
+from google.protobuf.internal import builder as _builder
+# @@protoc_insertion_point(imports)
+
+_sym_db = _symbol_database.Default()
+
+
+from flyteidl.core import interface_pb2 as flyteidl_dot_core_dot_interface__pb2
+
+
+DESCRIPTOR = _descriptor_pool.Default().AddSerializedFile(b'\n\x15workflow/common.proto\x12\x11\x63loudidl.workflow\x1a\x1d\x66lyteidl/core/interface.proto\"\\\n\x0eNamedParameter\x12\x12\n\x04name\x18\x01 \x01(\tR\x04name\x12\x36\n\tparameter\x18\x02 \x01(\x0b\x32\x18.flyteidl.core.ParameterR\tparameterB\xb8\x01\n\x15\x63om.cloudidl.workflowB\x0b\x43ommonProtoH\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)
+_builder.BuildTopDescriptorsAndMessages(DESCRIPTOR, 'workflow.common_pb2', _globals)
+if _descriptor._USE_C_DESCRIPTORS == False:
+  DESCRIPTOR._options = None
+  DESCRIPTOR._serialized_options = b'\n\025com.cloudidl.workflowB\013CommonProtoH\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'
+  _globals['_NAMEDPARAMETER']._serialized_start=75
+  _globals['_NAMEDPARAMETER']._serialized_end=167
+# @@protoc_insertion_point(module_scope)

flyte/_protos/workflow/common_pb2.pyi

@@ -0,0 +1,14 @@
+from flyteidl.core import interface_pb2 as _interface_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 NamedParameter(_message.Message):
+    __slots__ = ["name", "parameter"]
+    NAME_FIELD_NUMBER: _ClassVar[int]
+    PARAMETER_FIELD_NUMBER: _ClassVar[int]
+    name: str
+    parameter: _interface_pb2.Parameter
+    def __init__(self, name: _Optional[str] = ..., parameter: _Optional[_Union[_interface_pb2.Parameter, _Mapping]] = ...) -> None: ...

flyte/_protos/workflow/common_pb2_grpc.py

@@ -0,0 +1,4 @@
+# Generated by the gRPC Python protocol compiler plugin. DO NOT EDIT!
+"""Client and server classes corresponding to protobuf-defined services."""
+import grpc
+

flyte/_protos/workflow/queue_service_pb2.py

@@ -17,7 +17,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\x19\x66lyteidl/core/types.proto\x1a\x17validate/validate.proto\x1a\x1dworkflow/run_definition.proto\x1a\x1eworkflow/task_definition.proto\"\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\"f\n\x18\x41\x62ortQueuedActionRequest\x12J\n\taction_id\x18\x01 \x01(\x0b\x32#.cloudidl.workflow.ActionIdentifierB\x08\xfa\x42\x05\x8a\x01\x02\x10\x01R\x08\x61\x63tionId\"\x1b\n\x19\x41\x62ortQueuedActionResponse\"\x8e\x03\n\x10HeartbeatRequest\x12$\n\tworker_id\x18\x01 \x01(\tB\x07\xfa\x42\x04r\x02\x10\x01R\x08workerId\x12&\n\ncluster_id\x18\x02 \x01(\tB\x07\xfa\x42\x04r\x02\x10\x01R\tclusterId\x12O\n\x11\x61\x63tive_action_ids\x18\x03 \x03(\x0b\x32#.cloudidl.workflow.ActionIdentifierR\x0f\x61\x63tiveActionIds\x12S\n\x13terminal_action_ids\x18\x04 \x03(\x0b\x32#.cloudidl.workflow.ActionIdentifierR\x11terminalActionIds\x12-\n\x0corganization\x18\x05 \x01(\tB\x07\xfa\x42\x04r\x02\x10\x01H\x00R\x0corganization\x12-\n\x12\x61vailable_capacity\x18\x06 \x01(\x05R\x11\x61vailableCapacity\x12\x19\n\x03org\x18\x07 \x01(\tB\x07\xfa\x42\x04r\x02\x10\x01R\x03orgB\r\n\x06\x66ilter\x12\x03\xf8\x42\x01\"\xf6\x01\n\x11HeartbeatResponse\x12\x37\n\nnew_leases\x18\x01 \x03(\x0b\x32\x18.cloudidl.workflow.LeaseR\tnewLeases\x12Q\n\x12\x61\x62orted_action_ids\x18\x02 \x03(\x0b\x32#.cloudidl.workflow.ActionIdentifierR\x10\x61\x62ortedActionIds\x12U\n\x14\x66inalized_action_ids\x18\x03 \x03(\x0b\x32#.cloudidl.workflow.ActionIdentifierR\x12\x66inalizedActionIds\"V\n\x13StreamLeasesRequest\x12$\n\tworker_id\x18\x01 \x01(\tB\x07\xfa\x42\x04r\x02\x10\x01R\x08workerId\x12\x19\n\x03org\x18\x02 \x01(\tB\x07\xfa\x42\x04r\x02\x10\x01R\x03org\"H\n\x14StreamLeasesResponse\x12\x30\n\x06leases\x18\x01 \x03(\x0b\x32\x18.cloudidl.workflow.LeaseR\x06leases\"\xe5\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\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=\n\x04task\x18\x0b \x01(\x0b\x32\x1d.cloudidl.workflow.TaskActionB\x08\xfa\x42\x05\x8a\x01\x02\x10\x01H\x00R\x04task\x12L\n\tcondition\x18\x0c \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\x07subject\x12%\n\x0eprevious_state\x18\n \x01(\tR\rpreviousStateB\x0b\n\x04spec\x12\x03\xf8\x42\x01\x42\x15\n\x13_parent_action_name2\xa9\x03\n\x0cQueueService\x12\x64\n\rEnqueueAction\x12\'.cloudidl.workflow.EnqueueActionRequest\x1a(.cloudidl.workflow.EnqueueActionResponse\"\x00\x12p\n\x11\x41\x62ortQueuedAction\x12+.cloudidl.workflow.AbortQueuedActionRequest\x1a,.cloudidl.workflow.AbortQueuedActionResponse\"\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\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')
 
 _globals = globals()
 _builder.BuildMessageAndEnumDescriptors(DESCRIPTOR, _globals)
@@ -25,6 +25,12 @@ _builder.BuildTopDescriptorsAndMessages(DESCRIPTOR, 'workflow.queue_service_pb2'
 if _descriptor._USE_C_DESCRIPTORS == False:
   DESCRIPTOR._options = None
   DESCRIPTOR._serialized_options = b'\n\025com.cloudidl.workflowB\021QueueServiceProtoH\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'
+  _WORKERIDENTIFIER.fields_by_name['organization']._options = None
+  _WORKERIDENTIFIER.fields_by_name['organization']._serialized_options = b'\372B\004r\002\020\001'
+  _WORKERIDENTIFIER.fields_by_name['cluster']._options = None
+  _WORKERIDENTIFIER.fields_by_name['cluster']._serialized_options = b'\372B\004r\002\020\001'
+  _WORKERIDENTIFIER.fields_by_name['name']._options = None
+  _WORKERIDENTIFIER.fields_by_name['name']._serialized_options = b'\372B\004r\002\020\001'
   _ENQUEUEACTIONREQUEST.fields_by_name['action_id']._options = None
   _ENQUEUEACTIONREQUEST.fields_by_name['action_id']._serialized_options = b'\372B\005\212\001\002\020\001'
   _ENQUEUEACTIONREQUEST.fields_by_name['input_uri']._options = None
@@ -49,22 +55,12 @@ if _descriptor._USE_C_DESCRIPTORS == False:
   _CONDITIONACTION.fields_by_name['run_id']._serialized_options = b'\372B\004r\002\020\001'
   _CONDITIONACTION.fields_by_name['action_id']._options = None
   _CONDITIONACTION.fields_by_name['action_id']._serialized_options = b'\372B\004r\002\020\001'
-  _ABORTQUEUEDACTIONREQUEST.fields_by_name['action_id']._options = None
-  _ABORTQUEUEDACTIONREQUEST.fields_by_name['action_id']._serialized_options = b'\372B\005\212\001\002\020\001'
-  _HEARTBEATREQUEST.oneofs_by_name['filter']._options = None
-  _HEARTBEATREQUEST.oneofs_by_name['filter']._serialized_options = b'\370B\001'
+  _ABORTQUEUEDRUNREQUEST.fields_by_name['run_id']._options = None
+  _ABORTQUEUEDRUNREQUEST.fields_by_name['run_id']._serialized_options = b'\372B\005\212\001\002\020\001'
   _HEARTBEATREQUEST.fields_by_name['worker_id']._options = None
-  _HEARTBEATREQUEST.fields_by_name['worker_id']._serialized_options = b'\372B\004r\002\020\001'
-  _HEARTBEATREQUEST.fields_by_name['cluster_id']._options = None
-  _HEARTBEATREQUEST.fields_by_name['cluster_id']._serialized_options = b'\372B\004r\002\020\001'
-  _HEARTBEATREQUEST.fields_by_name['organization']._options = None
-  _HEARTBEATREQUEST.fields_by_name['organization']._serialized_options = b'\372B\004r\002\020\001'
-  _HEARTBEATREQUEST.fields_by_name['org']._options = None
-  _HEARTBEATREQUEST.fields_by_name['org']._serialized_options = b'\372B\004r\002\020\001'
+  _HEARTBEATREQUEST.fields_by_name['worker_id']._serialized_options = b'\372B\005\212\001\002\020\001'
   _STREAMLEASESREQUEST.fields_by_name['worker_id']._options = None
-  _STREAMLEASESREQUEST.fields_by_name['worker_id']._serialized_options = b'\372B\004r\002\020\001'
-  _STREAMLEASESREQUEST.fields_by_name['org']._options = None
-  _STREAMLEASESREQUEST.fields_by_name['org']._serialized_options = b'\372B\004r\002\020\001'
+  _STREAMLEASESREQUEST.fields_by_name['worker_id']._serialized_options = b'\372B\005\212\001\002\020\001'
   _LEASE.oneofs_by_name['spec']._options = None
   _LEASE.oneofs_by_name['spec']._serialized_options = b'\370B\001'
   _LEASE.fields_by_name['action_id']._options = None
@@ -77,30 +73,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['_ENQUEUEACTIONREQUEST']._serialized_start=167
-  _globals['_ENQUEUEACTIONREQUEST']._serialized_end=689
-  _globals['_TASKACTION']._serialized_start=691
-  _globals['_TASKACTION']._serialized_end=813
-  _globals['_TRACEACTION']._serialized_start=815
-  _globals['_TRACEACTION']._serialized_end=857
-  _globals['_CONDITIONACTION']._serialized_start=860
-  _globals['_CONDITIONACTION']._serialized_end=1126
-  _globals['_ENQUEUEACTIONRESPONSE']._serialized_start=1128
-  _globals['_ENQUEUEACTIONRESPONSE']._serialized_end=1151
-  _globals['_ABORTQUEUEDACTIONREQUEST']._serialized_start=1153
-  _globals['_ABORTQUEUEDACTIONREQUEST']._serialized_end=1255
-  _globals['_ABORTQUEUEDACTIONRESPONSE']._serialized_start=1257
-  _globals['_ABORTQUEUEDACTIONRESPONSE']._serialized_end=1284
-  _globals['_HEARTBEATREQUEST']._serialized_start=1287
-  _globals['_HEARTBEATREQUEST']._serialized_end=1685
-  _globals['_HEARTBEATRESPONSE']._serialized_start=1688
-  _globals['_HEARTBEATRESPONSE']._serialized_end=1934
-  _globals['_STREAMLEASESREQUEST']._serialized_start=1936
-  _globals['_STREAMLEASESREQUEST']._serialized_end=2022
-  _globals['_STREAMLEASESRESPONSE']._serialized_start=2024
-  _globals['_STREAMLEASESRESPONSE']._serialized_end=2096
-  _globals['_LEASE']._serialized_start=2099
-  _globals['_LEASE']._serialized_end=2584
-  _globals['_QUEUESERVICE']._serialized_start=2587
-  _globals['_QUEUESERVICE']._serialized_end=3012
+  _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
 # @@protoc_insertion_point(module_scope)

flyte/_protos/workflow/queue_service_pb2.pyi

@@ -9,6 +9,16 @@ from typing import ClassVar as _ClassVar, Iterable as _Iterable, Mapping as _Map
 
 DESCRIPTOR: _descriptor.FileDescriptor
 
+class WorkerIdentifier(_message.Message):
+    __slots__ = ["organization", "cluster", "name"]
+    ORGANIZATION_FIELD_NUMBER: _ClassVar[int]
+    CLUSTER_FIELD_NUMBER: _ClassVar[int]
+    NAME_FIELD_NUMBER: _ClassVar[int]
+    organization: str
+    cluster: str
+    name: str
+    def __init__(self, organization: _Optional[str] = ..., cluster: _Optional[str] = ..., name: _Optional[str] = ...) -> None: ...
+
 class EnqueueActionRequest(_message.Message):
     __slots__ = ["action_id", "parent_action_name", "input_uri", "run_output_base", "group", "subject", "task", "trace", "condition"]
     ACTION_ID_FIELD_NUMBER: _ClassVar[int]
@@ -66,51 +76,45 @@ class EnqueueActionResponse(_message.Message):
     __slots__ = []
     def __init__(self) -> None: ...
 
-class AbortQueuedActionRequest(_message.Message):
-    __slots__ = ["action_id"]
-    ACTION_ID_FIELD_NUMBER: _ClassVar[int]
-    action_id: _run_definition_pb2.ActionIdentifier
-    def __init__(self, action_id: _Optional[_Union[_run_definition_pb2.ActionIdentifier, _Mapping]] = ...) -> None: ...
+class AbortQueuedRunRequest(_message.Message):
+    __slots__ = ["run_id"]
+    RUN_ID_FIELD_NUMBER: _ClassVar[int]
+    run_id: _run_definition_pb2.RunIdentifier
+    def __init__(self, run_id: _Optional[_Union[_run_definition_pb2.RunIdentifier, _Mapping]] = ...) -> None: ...
 
-class AbortQueuedActionResponse(_message.Message):
+class AbortQueuedRunResponse(_message.Message):
     __slots__ = []
     def __init__(self) -> None: ...
 
 class HeartbeatRequest(_message.Message):
-    __slots__ = ["worker_id", "cluster_id", "active_action_ids", "terminal_action_ids", "organization", "available_capacity", "org"]
+    __slots__ = ["worker_id", "active_action_ids", "terminal_action_ids", "aborted_action_ids", "available_capacity"]
     WORKER_ID_FIELD_NUMBER: _ClassVar[int]
-    CLUSTER_ID_FIELD_NUMBER: _ClassVar[int]
     ACTIVE_ACTION_IDS_FIELD_NUMBER: _ClassVar[int]
     TERMINAL_ACTION_IDS_FIELD_NUMBER: _ClassVar[int]
-    ORGANIZATION_FIELD_NUMBER: _ClassVar[int]
+    ABORTED_ACTION_IDS_FIELD_NUMBER: _ClassVar[int]
     AVAILABLE_CAPACITY_FIELD_NUMBER: _ClassVar[int]
-    ORG_FIELD_NUMBER: _ClassVar[int]
-    worker_id: str
-    cluster_id: str
+    worker_id: WorkerIdentifier
     active_action_ids: _containers.RepeatedCompositeFieldContainer[_run_definition_pb2.ActionIdentifier]
     terminal_action_ids: _containers.RepeatedCompositeFieldContainer[_run_definition_pb2.ActionIdentifier]
-    organization: str
+    aborted_action_ids: _containers.RepeatedCompositeFieldContainer[_run_definition_pb2.ActionIdentifier]
     available_capacity: int
-    org: str
-    def __init__(self, worker_id: _Optional[str] = ..., cluster_id: _Optional[str] = ..., active_action_ids: _Optional[_Iterable[_Union[_run_definition_pb2.ActionIdentifier, _Mapping]]] = ..., terminal_action_ids: _Optional[_Iterable[_Union[_run_definition_pb2.ActionIdentifier, _Mapping]]] = ..., organization: _Optional[str] = ..., available_capacity: _Optional[int] = ..., org: _Optional[str] = ...) -> None: ...
+    def __init__(self, worker_id: _Optional[_Union[WorkerIdentifier, _Mapping]] = ..., active_action_ids: _Optional[_Iterable[_Union[_run_definition_pb2.ActionIdentifier, _Mapping]]] = ..., terminal_action_ids: _Optional[_Iterable[_Union[_run_definition_pb2.ActionIdentifier, _Mapping]]] = ..., aborted_action_ids: _Optional[_Iterable[_Union[_run_definition_pb2.ActionIdentifier, _Mapping]]] = ..., available_capacity: _Optional[int] = ...) -> None: ...
 
 class HeartbeatResponse(_message.Message):
-    __slots__ = ["new_leases", "aborted_action_ids", "finalized_action_ids"]
+    __slots__ = ["new_leases", "aborted_leases", "finalized_action_ids"]
     NEW_LEASES_FIELD_NUMBER: _ClassVar[int]
-    ABORTED_ACTION_IDS_FIELD_NUMBER: _ClassVar[int]
+    ABORTED_LEASES_FIELD_NUMBER: _ClassVar[int]
     FINALIZED_ACTION_IDS_FIELD_NUMBER: _ClassVar[int]
     new_leases: _containers.RepeatedCompositeFieldContainer[Lease]
-    aborted_action_ids: _containers.RepeatedCompositeFieldContainer[_run_definition_pb2.ActionIdentifier]
+    aborted_leases: _containers.RepeatedCompositeFieldContainer[Lease]
     finalized_action_ids: _containers.RepeatedCompositeFieldContainer[_run_definition_pb2.ActionIdentifier]
-    def __init__(self, new_leases: _Optional[_Iterable[_Union[Lease, _Mapping]]] = ..., aborted_action_ids: _Optional[_Iterable[_Union[_run_definition_pb2.ActionIdentifier, _Mapping]]] = ..., finalized_action_ids: _Optional[_Iterable[_Union[_run_definition_pb2.ActionIdentifier, _Mapping]]] = ...) -> None: ...
+    def __init__(self, new_leases: _Optional[_Iterable[_Union[Lease, _Mapping]]] = ..., aborted_leases: _Optional[_Iterable[_Union[Lease, _Mapping]]] = ..., finalized_action_ids: _Optional[_Iterable[_Union[_run_definition_pb2.ActionIdentifier, _Mapping]]] = ...) -> None: ...
 
 class StreamLeasesRequest(_message.Message):
-    __slots__ = ["worker_id", "org"]
+    __slots__ = ["worker_id"]
     WORKER_ID_FIELD_NUMBER: _ClassVar[int]
-    ORG_FIELD_NUMBER: _ClassVar[int]
-    worker_id: str
-    org: str
-    def __init__(self, worker_id: _Optional[str] = ..., org: _Optional[str] = ...) -> None: ...
+    worker_id: WorkerIdentifier
+    def __init__(self, worker_id: _Optional[_Union[WorkerIdentifier, _Mapping]] = ...) -> None: ...
 
 class StreamLeasesResponse(_message.Message):
     __slots__ = ["leases"]
@@ -119,7 +123,7 @@ class StreamLeasesResponse(_message.Message):
     def __init__(self, leases: _Optional[_Iterable[_Union[Lease, _Mapping]]] = ...) -> None: ...
 
 class Lease(_message.Message):
-    __slots__ = ["action_id", "parent_action_name", "input_uri", "run_output_base", "task", "condition", "group", "subject", "previous_state"]
+    __slots__ = ["action_id", "parent_action_name", "input_uri", "run_output_base", "task", "condition", "group", "subject"]
     ACTION_ID_FIELD_NUMBER: _ClassVar[int]
     PARENT_ACTION_NAME_FIELD_NUMBER: _ClassVar[int]
     INPUT_URI_FIELD_NUMBER: _ClassVar[int]
@@ -128,7 +132,6 @@ class Lease(_message.Message):
     CONDITION_FIELD_NUMBER: _ClassVar[int]
     GROUP_FIELD_NUMBER: _ClassVar[int]
     SUBJECT_FIELD_NUMBER: _ClassVar[int]
-    PREVIOUS_STATE_FIELD_NUMBER: _ClassVar[int]
     action_id: _run_definition_pb2.ActionIdentifier
     parent_action_name: str
     input_uri: str
@@ -137,5 +140,4 @@ class Lease(_message.Message):
     condition: ConditionAction
     group: str
     subject: str
-    previous_state: str
-    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] = ..., task: _Optional[_Union[TaskAction, _Mapping]] = ..., condition: _Optional[_Union[ConditionAction, _Mapping]] = ..., group: _Optional[str] = ..., subject: _Optional[str] = ..., previous_state: _Optional[str] = ...) -> None: ...
+    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] = ..., task: _Optional[_Union[TaskAction, _Mapping]] = ..., condition: _Optional[_Union[ConditionAction, _Mapping]] = ..., group: _Optional[str] = ..., subject: _Optional[str] = ...) -> None: ...