flyte 0.2.0b35__py3-none-any.whl → 0.2.0b37__py3-none-any.whl

flyte/_image.py

@@ -262,7 +262,7 @@ class Env(Layer):
 
 Architecture = Literal["linux/amd64", "linux/arm64"]
 
-_BASE_REGISTRY = "ghcr.io/unionai-oss"
+_BASE_REGISTRY = "ghcr.io/flyteorg"
 _DEFAULT_IMAGE_NAME = "flyte"
 
 

flyte/_internal/controllers/_local_controller.py

@@ -161,10 +161,10 @@ class LocalController:
         assert action_output_path
         return (
             TraceInfo(
+                name=_func.__name__,
                 action=action_id,
                 interface=_interface,
                 inputs_path=action_output_path,
-                name=_func.__name__,
             ),
             True,
         )
@@ -189,7 +189,8 @@ class LocalController:
             converted_error = convert.convert_from_native_to_error(info.error)
             assert converted_error
         assert info.action
-        assert info.duration
+        assert info.start_time
+        assert info.end_time
 
     async def submit_task_ref(self, _task: task_definition_pb2.TaskDetails, *args, **kwargs) -> Any:
         raise flyte.errors.ReferenceTaskError("Reference tasks cannot be executed locally, only remotely.")

flyte/_internal/controllers/_trace.py

@@ -1,5 +1,4 @@
-from dataclasses import dataclass
-from datetime import timedelta
+from dataclasses import dataclass, field
 from typing import Any, Optional
 
 from flyte.models import ActionID, NativeInterface
@@ -12,30 +11,35 @@ class TraceInfo:
      the action is completed.
     """
 
+    name: str
     action: ActionID
     interface: NativeInterface
     inputs_path: str
-    duration: Optional[timedelta] = None
+    start_time: float = field(init=False, default=0.0)
+    end_time: float = field(init=False, default=0.0)
     output: Optional[Any] = None
     error: Optional[Exception] = None
-    name: str = ""
 
-    def add_outputs(self, output: Any, duration: timedelta):
+    def add_outputs(self, output: Any, start_time: float, end_time: float):
         """
         Add outputs to the trace information.
         :param output: Output of the action
-        :param duration: Duration of the action
+        :param start_time: Start time of the action
+        :param end_time: End time of the action
         :return:
         """
         self.output = output
-        self.duration = duration
+        self.start_time = start_time
+        self.end_time = end_time
 
-    def add_error(self, error: Exception, duration: timedelta):
+    def add_error(self, error: Exception, start_time: float, end_time: float):
         """
         Add error to the trace information.
         :param error: Error of the action
-        :param duration: Duration of the action
+        :param start_time: Start time of the action
+        :param end_time: End time of the action
         :return:
         """
         self.error = error
-        self.duration = duration
+        self.start_time = start_time
+        self.end_time = end_time

flyte/_internal/controllers/remote/_action.py

@@ -1,11 +1,18 @@
 from __future__ import annotations
 
 from dataclasses import dataclass
-from typing import Any, Literal
+from typing import Literal
 
 from flyteidl.core import execution_pb2
-
-from flyte._protos.workflow import run_definition_pb2, state_service_pb2, task_definition_pb2
+from google.protobuf import timestamp_pb2
+
+from flyte._protos.common import identifier_pb2
+from flyte._protos.workflow import (
+    queue_service_pb2,
+    run_definition_pb2,
+    state_service_pb2,
+    task_definition_pb2,
+)
 from flyte.models import GroupData
 
 ActionType = Literal["task", "trace"]
@@ -18,12 +25,13 @@ class Action:
     Holds the inmemory state of a task. It is combined representation of local and remote states.
     """
 
-    action_id: run_definition_pb2.ActionIdentifier
+    action_id: identifier_pb2.ActionIdentifier
     parent_action_name: str
     type: ActionType = "task"  # type of action, task or trace
     friendly_name: str | None = None
     group: GroupData | None = None
     task: task_definition_pb2.TaskSpec | None = None
+    trace: queue_service_pb2.TraceAction | None = None
     inputs_uri: str | None = None
     run_output_base: str | None = None
     realized_outputs_uri: str | None = None
@@ -108,7 +116,7 @@ class Action:
     def from_task(
         cls,
         parent_action_name: str,
-        sub_action_id: run_definition_pb2.ActionIdentifier,
+        sub_action_id: identifier_pb2.ActionIdentifier,
         group_data: GroupData | None,
         task_spec: task_definition_pb2.TaskSpec,
         inputs_uri: str,
@@ -153,16 +161,27 @@ class Action:
     def from_trace(
         cls,
         parent_action_name: str,
-        action_id: run_definition_pb2.ActionIdentifier,
+        action_id: identifier_pb2.ActionIdentifier,
         friendly_name: str,
         group_data: GroupData | None,
-        trace_spec: Any,  # TODO
         inputs_uri: str,
         outputs_uri: str,
+        start_time: float,  # Unix timestamp in seconds with fractional seconds
+        end_time: float,  # Unix timestamp in seconds with fractional seconds
+        run_output_base: str,
+        report_uri: str | None = None,
     ) -> Action:
         """
         This creates a new action for tracing purposes. It is used to track the execution of a trace.
         """
+        st = timestamp_pb2.Timestamp()
+        st.FromSeconds(int(start_time))
+        st.nanos = int((start_time % 1) * 1e9)
+
+        et = timestamp_pb2.Timestamp()
+        et.FromSeconds(int(end_time))
+        et.nanos = int((end_time % 1) * 1e9)
+
         return cls(
             action_id=action_id,
             parent_action_name=parent_action_name,
@@ -172,4 +191,15 @@ class Action:
             inputs_uri=inputs_uri,
             realized_outputs_uri=outputs_uri,
             phase=run_definition_pb2.Phase.PHASE_SUCCEEDED,
+            run_output_base=run_output_base,
+            trace=queue_service_pb2.TraceAction(
+                name=friendly_name,
+                phase=run_definition_pb2.Phase.PHASE_SUCCEEDED,
+                start_time=st,
+                end_time=et,
+                outputs=run_definition_pb2.OutputReferences(
+                    output_uri=outputs_uri,
+                    report_uri=report_uri,
+                ),
+            ),
         )

flyte/_internal/controllers/remote/_controller.py

@@ -22,6 +22,7 @@ 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._logging import logger
+from flyte._protos.common import identifier_pb2
 from flyte._protos.workflow import run_definition_pb2, task_definition_pb2
 from flyte._task import TaskTemplate
 from flyte._utils.helpers import _selector_policy
@@ -131,11 +132,18 @@ class RemoteController(Controller):
         Generate a task call sequence for the given task object and action ID.
         This is used to track the number of times a task is called within an action.
         """
-        current_action_sequencer = self._parent_action_task_call_sequence[unique_action_name(action_id)]
+        uniq = unique_action_name(action_id)
+        current_action_sequencer = self._parent_action_task_call_sequence[uniq]
         current_task_id = id(task_obj)
         v = current_action_sequencer[current_task_id]
         new_seq = v + 1
         current_action_sequencer[current_task_id] = new_seq
+        name = ""
+        if hasattr(task_obj, "__name__"):
+            name = task_obj.__name__
+        elif hasattr(task_obj, "name"):
+            name = task_obj.name
+        logger.warning(f"For action {uniq}, task {name} call sequence is {new_seq}")
         return new_seq
 
     async def _submit(self, _task_call_seq: int, _task: TaskTemplate, *args, **kwargs) -> Any:
@@ -178,6 +186,7 @@ class RemoteController(Controller):
         sub_action_id, sub_action_output_path = convert.generate_sub_action_id_and_output_path(
             tctx, task_spec, inputs_hash, _task_call_seq
         )
+        logger.warning(f"Sub action {sub_action_id} output path {sub_action_output_path}")
 
         serialized_inputs = inputs.proto_inputs.SerializeToString(deterministic=True)
         inputs_uri = io.inputs_path(sub_action_output_path)
@@ -204,9 +213,9 @@ class RemoteController(Controller):
         inputs_hash = None  # type: ignore
 
         action = Action.from_task(
-            sub_action_id=run_definition_pb2.ActionIdentifier(
+            sub_action_id=identifier_pb2.ActionIdentifier(
                 name=sub_action_id.name,
-                run=run_definition_pb2.RunIdentifier(
+                run=identifier_pb2.RunIdentifier(
                     name=current_action_id.run_name,
                     project=current_action_id.project,
                     domain=current_action_id.domain,
@@ -293,7 +302,9 @@ class RemoteController(Controller):
                 self._submit_loop.set_exception_handler(exc_handler)
 
             self._submit_thread = threading.Thread(
-                name=f"remote-controller-{os.getpid()}-submitter", daemon=True, target=self._sync_thread_loop_runner
+                name=f"remote-controller-{os.getpid()}-submitter",
+                daemon=True,
+                target=self._sync_thread_loop_runner,
             )
             self._submit_thread.start()
 
@@ -307,7 +318,7 @@ class RemoteController(Controller):
         This method is invoked when the parent action is finished. It will finalize the run and upload the outputs
         to the control plane.
         """
-        run_id = run_definition_pb2.RunIdentifier(
+        run_id = identifier_pb2.RunIdentifier(
             name=action_id.run_name,
             project=action_id.project,
             domain=action_id.domain,
@@ -350,9 +361,9 @@ class RemoteController(Controller):
         serialized_inputs = None  # type: ignore
 
         prev_action = await self.get_action(
-            run_definition_pb2.ActionIdentifier(
+            identifier_pb2.ActionIdentifier(
                 name=sub_action_id.name,
-                run=run_definition_pb2.RunIdentifier(
+                run=identifier_pb2.RunIdentifier(
                     name=current_action_id.run_name,
                     project=current_action_id.project,
                     domain=current_action_id.domain,
@@ -363,21 +374,27 @@ class RemoteController(Controller):
         )
 
         if prev_action is None:
-            return TraceInfo(sub_action_id, _interface, inputs_uri), False
+            return TraceInfo(func_name, sub_action_id, _interface, inputs_uri), False
 
         if prev_action.phase == run_definition_pb2.PHASE_FAILED:
             if prev_action.has_error():
                 exc = convert.convert_error_to_native(prev_action.err)
-                return TraceInfo(sub_action_id, _interface, inputs_uri, error=exc), True
+                return (
+                    TraceInfo(func_name, sub_action_id, _interface, inputs_uri, error=exc),
+                    True,
+                )
             else:
                 logger.warning(f"Action {prev_action.action_id.name} failed, but no error was found, re-running trace!")
         elif prev_action.realized_outputs_uri is not None:
             outputs_file_path = io.outputs_path(prev_action.realized_outputs_uri)
             o = await io.load_outputs(outputs_file_path)
             outputs = await convert.convert_outputs_to_native(_interface, o)
-            return TraceInfo(sub_action_id, _interface, inputs_uri, output=outputs), True
+            return (
+                TraceInfo(func_name, sub_action_id, _interface, inputs_uri, output=outputs),
+                True,
+            )
 
-        return TraceInfo(sub_action_id, _interface, inputs_uri), False
+        return TraceInfo(func_name, sub_action_id, _interface, inputs_uri), False
 
     async def record_trace(self, info: TraceInfo):
         """
@@ -391,26 +408,29 @@ class RemoteController(Controller):
             raise flyte.errors.RuntimeSystemError("BadContext", "Task context not initialized")
 
         current_action_id = tctx.action
-        current_output_path = tctx.output_path
-        sub_run_output_path = storage.join(current_output_path, info.action.name)
+        sub_run_output_path = storage.join(tctx.run_base_dir, info.action.name)
+        print(f"Sub run output path for {info.name} is {sub_run_output_path}", flush=True)
 
         if info.interface.has_outputs():
             outputs_file_path: str = ""
             if info.output:
                 outputs = await convert.convert_from_native_to_outputs(info.output, info.interface)
                 outputs_file_path = io.outputs_path(sub_run_output_path)
-                await io.upload_outputs(outputs, outputs_file_path)
+                print(
+                    f"Uploading outputs for {info.name} Outputs file path: {outputs_file_path}",
+                    flush=True,
+                )
+                await io.upload_outputs(outputs, sub_run_output_path)
             elif info.error:
                 err = convert.convert_from_native_to_error(info.error)
-                error_path = io.error_path(sub_run_output_path)
-                await io.upload_error(err.err, error_path)
+                await io.upload_error(err.err, sub_run_output_path)
             else:
                 raise flyte.errors.RuntimeSystemError("BadTraceInfo", "Trace info does not have output or error")
             trace_action = Action.from_trace(
                 parent_action_name=current_action_id.name,
-                action_id=run_definition_pb2.ActionIdentifier(
+                action_id=identifier_pb2.ActionIdentifier(
                     name=info.action.name,
-                    run=run_definition_pb2.RunIdentifier(
+                    run=identifier_pb2.RunIdentifier(
                         name=current_action_id.run_name,
                         project=current_action_id.project,
                         domain=current_action_id.domain,
@@ -419,9 +439,11 @@ class RemoteController(Controller):
                 ),
                 inputs_uri=info.inputs_path,
                 outputs_uri=outputs_file_path,
-                trace_spec="",
                 friendly_name=info.name,
                 group_data=tctx.group_data,
+                run_output_base=tctx.run_base_dir,
+                start_time=info.start_time,
+                end_time=info.end_time,
             )
             try:
                 logger.info(
@@ -478,9 +500,9 @@ class RemoteController(Controller):
         inputs_hash = None  # type: ignore
 
         action = Action.from_task(
-            sub_action_id=run_definition_pb2.ActionIdentifier(
+            sub_action_id=identifier_pb2.ActionIdentifier(
                 name=sub_action_id.name,
-                run=run_definition_pb2.RunIdentifier(
+                run=identifier_pb2.RunIdentifier(
                     name=current_action_id.run_name,
                     project=current_action_id.project,
                     domain=current_action_id.domain,

flyte/_internal/controllers/remote/_core.py

@@ -11,7 +11,11 @@ from google.protobuf.wrappers_pb2 import StringValue
 
 import flyte.errors
 from flyte._logging import log, logger
-from flyte._protos.workflow import queue_service_pb2, run_definition_pb2, task_definition_pb2
+from flyte._protos.common import identifier_pb2
+from flyte._protos.workflow import (
+    queue_service_pb2,
+    task_definition_pb2,
+)
 from flyte.errors import RuntimeSystemError
 
 from ._action import Action
@@ -80,9 +84,7 @@ class Controller:
         """Public API to submit a resource and wait for completion"""
         return await self._run_coroutine_in_controller_thread(self._bg_submit_action(action))
 
-    async def get_action(
-        self, action_id: run_definition_pb2.ActionIdentifier, parent_action_name: str
-    ) -> Optional[Action]:
+    async def get_action(self, action_id: identifier_pb2.ActionIdentifier, parent_action_name: str) -> Optional[Action]:
         """Get the action from the informer"""
         informer = await self._informers.get(run_name=action_id.run.name, parent_action_name=parent_action_name)
         if informer:
@@ -94,7 +96,10 @@ class Controller:
         return await self._run_coroutine_in_controller_thread(self._bg_cancel_action(action))
 
     async def _finalize_parent_action(
-        self, run_id: run_definition_pb2.RunIdentifier, parent_action_name: str, timeout: Optional[float] = None
+        self,
+        run_id: identifier_pb2.RunIdentifier,
+        parent_action_name: str,
+        timeout: Optional[float] = None,
     ):
         """Finalize the parent run"""
         await self._run_coroutine_in_controller_thread(
@@ -124,7 +129,8 @@ class Controller:
         """Watch for errors in the background thread"""
         await self._run_coroutine_in_controller_thread(self._bg_watch_for_errors())
         raise RuntimeSystemError(
-            code="InformerWatchFailure", message=f"Controller thread failed with exception: {self._get_exception()}"
+            code="InformerWatchFailure",
+            message=f"Controller thread failed with exception: {self._get_exception()}",
         )
 
     @log
@@ -162,7 +168,8 @@ class Controller:
 
         if self._get_exception():
             raise RuntimeSystemError(
-                type(self._get_exception()).__name__, f"Controller thread startup failed: {self._get_exception()}"
+                type(self._get_exception()).__name__,
+                f"Controller thread startup failed: {self._get_exception()}",
             )
 
         logger.info(f"Controller started in thread: {self._thread.name}")
@@ -223,7 +230,7 @@ class Controller:
             logger.debug(f"Controller thread exiting: {threading.current_thread().name}")
 
     async def _bg_get_action(
-        self, action_id: run_definition_pb2.ActionIdentifier, parent_action_name: str
+        self, action_id: identifier_pb2.ActionIdentifier, parent_action_name: str
     ) -> Optional[Action]:
         """Get the action from the informer"""
         # Ensure the informer is created and wait for it to be ready
@@ -240,7 +247,10 @@ class Controller:
         return None
 
     async def _bg_finalize_informer(
-        self, run_id: run_definition_pb2.RunIdentifier, parent_action_name: str, timeout: Optional[float] = None
+        self,
+        run_id: identifier_pb2.RunIdentifier,
+        parent_action_name: str,
+        timeout: Optional[float] = None,
     ):
         informer = await self._informers.remove(run_name=run_id.name, parent_action_name=parent_action_name)
         if informer:
@@ -294,7 +304,10 @@ class Controller:
                 # )
                 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]:
+                if e.code() in [
+                    grpc.StatusCode.NOT_FOUND,
+                    grpc.StatusCode.FAILED_PRECONDITION,
+                ]:
                     logger.info(f"Action {action.name} not found, assumed completed or cancelled.")
                     return
         else:
@@ -333,10 +346,8 @@ class Controller:
                     spec=action.task,
                     cache_key=cache_key,
                 )
-            else:
-                trace = queue_service_pb2.TraceAction(
-                    name=action.friendly_name,
-                )
+            elif action.type == "trace":
+                trace = action.trace
 
             logger.debug(f"Attempting to launch action: {action.name}")
             try:
@@ -380,7 +391,11 @@ class Controller:
     async def _bg_log_stats(self):
         """Periodically log resource stats if debug is enabled"""
         while self._running:
-            async for started, pending, terminal in self._informers.count_started_pending_terminal_actions():
+            async for (
+                started,
+                pending,
+                terminal,
+            ) in self._informers.count_started_pending_terminal_actions():
                 logger.info(f"Resource stats: Started={started}, Pending={pending}, Terminal={terminal}")
             await asyncio.sleep(self._resource_log_interval)
 
@@ -407,7 +422,8 @@ class Controller:
                     err.__cause__ = e
                     action.set_client_error(err)
                     informer = await self._informers.get(
-                        run_name=action.run_name, parent_action_name=action.parent_action_name
+                        run_name=action.run_name,
+                        parent_action_name=action.parent_action_name,
                     )
                     if informer:
                         await informer.fire_completion_event(action.name)

flyte/_internal/controllers/remote/_informer.py

@@ -7,6 +7,7 @@ from typing import AsyncIterator, Callable, Dict, Optional, Tuple, cast
 import grpc.aio
 
 from flyte._logging import log, logger
+from flyte._protos.common import identifier_pb2
 from flyte._protos.workflow import run_definition_pb2, state_service_pb2
 
 from ._action import Action
@@ -42,7 +43,9 @@ class ActionCache:
         if state.output_uri:
             logger.debug(f"Output URI: {state.output_uri}")
         else:
-            logger.warning(f"{state.action_id.name} has no output URI")
+            logger.warning(
+                f"{state.action_id.name} has no output URI, in phase {run_definition_pb2.Phase.Name(state.phase)}"
+            )
         if state.phase == run_definition_pb2.Phase.PHASE_FAILED:
             logger.error(
                 f"Action {state.action_id.name} failed with error (msg):"
@@ -125,7 +128,7 @@ class Informer:
 
     def __init__(
         self,
-        run_id: run_definition_pb2.RunIdentifier,
+        run_id: identifier_pb2.RunIdentifier,
         parent_action_name: str,
         shared_queue: Queue,
         client: Optional[StateService] = None,
@@ -217,7 +220,7 @@ class Informer:
             try:
                 watcher = self._client.Watch(
                     state_service_pb2.WatchRequest(
-                        parent_action_id=run_definition_pb2.ActionIdentifier(
+                        parent_action_id=identifier_pb2.ActionIdentifier(
                             name=self.parent_action_name,
                             run=self._run_id,
                         ),
@@ -288,7 +291,7 @@ class InformerCache:
     @log
     async def get_or_create(
         self,
-        run_id: run_definition_pb2.RunIdentifier,
+        run_id: identifier_pb2.RunIdentifier,
         parent_action_name: str,
         shared_queue: Queue,
         state_service: StateService,
@@ -330,20 +333,28 @@ class InformerCache:
     async def get(self, *, run_name: str, parent_action_name: str) -> Informer | None:
         """Get an informer by name"""
         async with self._lock:
-            return self._cache.get(Informer.mkname(run_name=run_name, parent_action_name=parent_action_name), None)
+            return self._cache.get(
+                Informer.mkname(run_name=run_name, parent_action_name=parent_action_name),
+                None,
+            )
 
     @log
     async def remove(self, *, run_name: str, parent_action_name: str) -> Informer | None:
         """Remove an informer from the cache"""
         async with self._lock:
-            return self._cache.pop(Informer.mkname(run_name=run_name, parent_action_name=parent_action_name), None)
+            return self._cache.pop(
+                Informer.mkname(run_name=run_name, parent_action_name=parent_action_name),
+                None,
+            )
 
     async def has(self, *, run_name: str, parent_action_name: str) -> bool:
         """Check if an informer exists in the cache"""
         async with self._lock:
             return Informer.mkname(run_name=run_name, parent_action_name=parent_action_name) in self._cache
 
-    async def count_started_pending_terminal_actions(self) -> AsyncIterator[Tuple[int, int, int]]:
+    async def count_started_pending_terminal_actions(
+        self,
+    ) -> AsyncIterator[Tuple[int, int, int]]:
         """Log resource stats"""
         async with self._lock:
             for informer in self._cache.values():

flyte/_internal/runtime/task_serde.py

@@ -59,7 +59,9 @@ def translate_task_to_wire(
     )
 
 
-def get_security_context(secrets: Optional[SecretRequest]) -> Optional[security_pb2.SecurityContext]:
+def get_security_context(
+    secrets: Optional[SecretRequest],
+) -> Optional[security_pb2.SecurityContext]:
     """
     Get the security context from a list of secrets. This is a placeholder function.
 
@@ -86,7 +88,9 @@ def get_security_context(secrets: Optional[SecretRequest]) -> Optional[security_
     )
 
 
-def get_proto_retry_strategy(retries: RetryStrategy | int | None) -> Optional[literals_pb2.RetryStrategy]:
+def get_proto_retry_strategy(
+    retries: RetryStrategy | int | None,
+) -> Optional[literals_pb2.RetryStrategy]:
     if retries is None:
         return None
 
@@ -158,7 +162,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.get_ignored_inputs() if cache_enabled else None,
+            cache_ignore_input_vars=(task_cache.get_ignored_inputs() if cache_enabled else None),
             runtime=tasks_pb2.RuntimeMetadata(
                 version=flyte.version(),
                 type=tasks_pb2.RuntimeMetadata.RuntimeType.FLYTE_SDK,
@@ -166,7 +170,7 @@ def get_proto_task(task: TaskTemplate, serialize_context: SerializationContext)
             ),
             retries=get_proto_retry_strategy(task.retries),
             timeout=get_proto_timeout(task.timeout),
-            pod_template_name=task.pod_template if task.pod_template and isinstance(task.pod_template, str) else None,
+            pod_template_name=(task.pod_template if task.pod_template and isinstance(task.pod_template, str) else None),
             interruptible=task.interruptable,
             generates_deck=wrappers_pb2.BoolValue(value=task.report),
         ),
@@ -300,7 +304,9 @@ def _get_k8s_pod(primary_container: tasks_pb2.Container, pod_template: PodTempla
     return tasks_pb2.K8sPod(pod_spec=pod_spec, metadata=metadata)
 
 
-def extract_code_bundle(task_spec: task_definition_pb2.TaskSpec) -> Optional[CodeBundle]:
+def extract_code_bundle(
+    task_spec: task_definition_pb2.TaskSpec,
+) -> Optional[CodeBundle]:
     """
     Extract the code bundle from the task spec.
     :param task_spec: The task spec to extract the code bundle from.
@@ -326,5 +332,10 @@ def extract_code_bundle(task_spec: task_definition_pb2.TaskSpec) -> Optional[Cod
                 # Extract the version from the argument
                 version = container.args[i + 1] if i + 1 < len(container.args) else ""
         if pkl_path or tgz_path:
-            return CodeBundle(destination=dest_path, tgz=tgz_path, pkl=pkl_path, computed_version=version)
+            return CodeBundle(
+                destination=dest_path,
+                tgz=tgz_path,
+                pkl=pkl_path,
+                computed_version=version,
+            )
     return None

flyte/_protos/common/identifier_pb2.py

@@ -14,7 +14,7 @@ _sym_db = _symbol_database.Default()
 from flyte._protos.validate.validate import validate_pb2 as validate_dot_validate__pb2
 
 
-DESCRIPTOR = _descriptor_pool.Default().AddSerializedFile(b'\n\x17\x63ommon/identifier.proto\x12\x0f\x63loudidl.common\x1a\x17validate/validate.proto\"~\n\x11ProjectIdentifier\x12+\n\x0corganization\x18\x01 \x01(\tB\x07\xfa\x42\x04r\x02\x10\x01R\x0corganization\x12\x1f\n\x06\x64omain\x18\x02 \x01(\tB\x07\xfa\x42\x04r\x02\x10\x01R\x06\x64omain\x12\x1b\n\x04name\x18\x03 \x01(\tB\x07\xfa\x42\x04r\x02\x10\x01R\x04name\"T\n\x11\x43lusterIdentifier\x12\"\n\x0corganization\x18\x01 \x01(\tR\x0corganization\x12\x1b\n\x04name\x18\x02 \x01(\tB\x07\xfa\x42\x04r\x02\x10\x01R\x04name\"O\n\x15\x43lusterPoolIdentifier\x12\"\n\x0corganization\x18\x01 \x01(\tR\x0corganization\x12\x12\n\x04name\x18\x02 \x01(\tR\x04name\"_\n\x17\x43lusterConfigIdentifier\x12+\n\x0corganization\x18\x01 \x01(\tB\x07\xfa\x42\x04r\x02\x10\x01R\x0corganization\x12\x17\n\x02id\x18\x02 \x01(\tB\x07\xfa\x42\x04r\x02\x10\x01R\x02id\"3\n\x0eUserIdentifier\x12!\n\x07subject\x18\x01 \x01(\tB\x07\xfa\x42\x04r\x02\x10\x01R\x07subject\":\n\x15\x41pplicationIdentifier\x12!\n\x07subject\x18\x01 \x01(\tB\x07\xfa\x42\x04r\x02\x10\x01R\x07subject\"Q\n\x0eRoleIdentifier\x12\"\n\x0corganization\x18\x01 \x01(\tR\x0corganization\x12\x1b\n\x04name\x18\x02 \x01(\tB\x07\xfa\x42\x04r\x02\x10\x01R\x04name\"O\n\rOrgIdentifier\x12>\n\x04name\x18\x01 \x01(\tB*\xfa\x42\'r%\x10\x01\x18?2\x1f^[a-z0-9]([-a-z0-9]*[a-z0-9])?$R\x04name\"y\n\x18ManagedClusterIdentifier\x12\x1b\n\x04name\x18\x02 \x01(\tB\x07\xfa\x42\x04r\x02\x10\x01R\x04name\x12:\n\x03org\x18\x03 \x01(\x0b\x32\x1e.cloudidl.common.OrgIdentifierB\x08\xfa\x42\x05\x8a\x01\x02\x10\x01R\x03orgJ\x04\x08\x01\x10\x02\"S\n\x10PolicyIdentifier\x12\"\n\x0corganization\x18\x01 \x01(\tR\x0corganization\x12\x1b\n\x04name\x18\x02 \x01(\tB\x07\xfa\x42\x04r\x02\x10\x01R\x04nameB\xb0\x01\n\x13\x63om.cloudidl.commonB\x0fIdentifierProtoH\x02P\x01Z)github.com/unionai/cloud/gen/pb-go/common\xa2\x02\x03\x43\x43X\xaa\x02\x0f\x43loudidl.Common\xca\x02\x0f\x43loudidl\\Common\xe2\x02\x1b\x43loudidl\\Common\\GPBMetadata\xea\x02\x10\x43loudidl::Commonb\x06proto3')
+DESCRIPTOR = _descriptor_pool.Default().AddSerializedFile(b'\n\x17\x63ommon/identifier.proto\x12\x0f\x63loudidl.common\x1a\x17validate/validate.proto\"~\n\x11ProjectIdentifier\x12+\n\x0corganization\x18\x01 \x01(\tB\x07\xfa\x42\x04r\x02\x10\x01R\x0corganization\x12\x1f\n\x06\x64omain\x18\x02 \x01(\tB\x07\xfa\x42\x04r\x02\x10\x01R\x06\x64omain\x12\x1b\n\x04name\x18\x03 \x01(\tB\x07\xfa\x42\x04r\x02\x10\x01R\x04name\"T\n\x11\x43lusterIdentifier\x12\"\n\x0corganization\x18\x01 \x01(\tR\x0corganization\x12\x1b\n\x04name\x18\x02 \x01(\tB\x07\xfa\x42\x04r\x02\x10\x01R\x04name\"O\n\x15\x43lusterPoolIdentifier\x12\"\n\x0corganization\x18\x01 \x01(\tR\x0corganization\x12\x12\n\x04name\x18\x02 \x01(\tR\x04name\"_\n\x17\x43lusterConfigIdentifier\x12+\n\x0corganization\x18\x01 \x01(\tB\x07\xfa\x42\x04r\x02\x10\x01R\x0corganization\x12\x17\n\x02id\x18\x02 \x01(\tB\x07\xfa\x42\x04r\x02\x10\x01R\x02id\"3\n\x0eUserIdentifier\x12!\n\x07subject\x18\x01 \x01(\tB\x07\xfa\x42\x04r\x02\x10\x01R\x07subject\":\n\x15\x41pplicationIdentifier\x12!\n\x07subject\x18\x01 \x01(\tB\x07\xfa\x42\x04r\x02\x10\x01R\x07subject\"Q\n\x0eRoleIdentifier\x12\"\n\x0corganization\x18\x01 \x01(\tR\x0corganization\x12\x1b\n\x04name\x18\x02 \x01(\tB\x07\xfa\x42\x04r\x02\x10\x01R\x04name\"O\n\rOrgIdentifier\x12>\n\x04name\x18\x01 \x01(\tB*\xfa\x42\'r%\x10\x01\x18?2\x1f^[a-z0-9]([-a-z0-9]*[a-z0-9])?$R\x04name\"y\n\x18ManagedClusterIdentifier\x12\x1b\n\x04name\x18\x02 \x01(\tB\x07\xfa\x42\x04r\x02\x10\x01R\x04name\x12:\n\x03org\x18\x03 \x01(\x0b\x32\x1e.cloudidl.common.OrgIdentifierB\x08\xfa\x42\x05\x8a\x01\x02\x10\x01R\x03orgJ\x04\x08\x01\x10\x02\"S\n\x10PolicyIdentifier\x12\"\n\x0corganization\x18\x01 \x01(\tR\x0corganization\x12\x1b\n\x04name\x18\x02 \x01(\tB\x07\xfa\x42\x04r\x02\x10\x01R\x04name\"\x93\x01\n\rRunIdentifier\x12\x1b\n\x03org\x18\x01 \x01(\tB\t\xfa\x42\x06r\x04\x10\x01\x18?R\x03org\x12#\n\x07project\x18\x02 \x01(\tB\t\xfa\x42\x06r\x04\x10\x01\x18?R\x07project\x12!\n\x06\x64omain\x18\x03 \x01(\tB\t\xfa\x42\x06r\x04\x10\x01\x18?R\x06\x64omain\x12\x1d\n\x04name\x18\x04 \x01(\tB\t\xfa\x42\x06r\x04\x10\x01\x18\x1eR\x04name\"m\n\x10\x41\x63tionIdentifier\x12:\n\x03run\x18\x01 \x01(\x0b\x32\x1e.cloudidl.common.RunIdentifierB\x08\xfa\x42\x05\x8a\x01\x02\x10\x01R\x03run\x12\x1d\n\x04name\x18\x02 \x01(\tB\t\xfa\x42\x06r\x04\x10\x01\x18\x1eR\x04name\"\x86\x01\n\x17\x41\x63tionAttemptIdentifier\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!\n\x07\x61ttempt\x18\x02 \x01(\rB\x07\xfa\x42\x04*\x02 \x00R\x07\x61ttemptB\xb0\x01\n\x13\x63om.cloudidl.commonB\x0fIdentifierProtoH\x02P\x01Z)github.com/unionai/cloud/gen/pb-go/common\xa2\x02\x03\x43\x43X\xaa\x02\x0f\x43loudidl.Common\xca\x02\x0f\x43loudidl\\Common\xe2\x02\x1b\x43loudidl\\Common\\GPBMetadata\xea\x02\x10\x43loudidl::Commonb\x06proto3')
 
 _globals = globals()
 _builder.BuildMessageAndEnumDescriptors(DESCRIPTOR, _globals)
@@ -48,6 +48,22 @@ if _descriptor._USE_C_DESCRIPTORS == False:
   _MANAGEDCLUSTERIDENTIFIER.fields_by_name['org']._serialized_options = b'\372B\005\212\001\002\020\001'
   _POLICYIDENTIFIER.fields_by_name['name']._options = None
   _POLICYIDENTIFIER.fields_by_name['name']._serialized_options = b'\372B\004r\002\020\001'
+  _RUNIDENTIFIER.fields_by_name['org']._options = None
+  _RUNIDENTIFIER.fields_by_name['org']._serialized_options = b'\372B\006r\004\020\001\030?'
+  _RUNIDENTIFIER.fields_by_name['project']._options = None
+  _RUNIDENTIFIER.fields_by_name['project']._serialized_options = b'\372B\006r\004\020\001\030?'
+  _RUNIDENTIFIER.fields_by_name['domain']._options = None
+  _RUNIDENTIFIER.fields_by_name['domain']._serialized_options = b'\372B\006r\004\020\001\030?'
+  _RUNIDENTIFIER.fields_by_name['name']._options = None
+  _RUNIDENTIFIER.fields_by_name['name']._serialized_options = b'\372B\006r\004\020\001\030\036'
+  _ACTIONIDENTIFIER.fields_by_name['run']._options = None
+  _ACTIONIDENTIFIER.fields_by_name['run']._serialized_options = b'\372B\005\212\001\002\020\001'
+  _ACTIONIDENTIFIER.fields_by_name['name']._options = None
+  _ACTIONIDENTIFIER.fields_by_name['name']._serialized_options = b'\372B\006r\004\020\001\030\036'
+  _ACTIONATTEMPTIDENTIFIER.fields_by_name['action_id']._options = None
+  _ACTIONATTEMPTIDENTIFIER.fields_by_name['action_id']._serialized_options = b'\372B\005\212\001\002\020\001'
+  _ACTIONATTEMPTIDENTIFIER.fields_by_name['attempt']._options = None
+  _ACTIONATTEMPTIDENTIFIER.fields_by_name['attempt']._serialized_options = b'\372B\004*\002 \000'
   _globals['_PROJECTIDENTIFIER']._serialized_start=69
   _globals['_PROJECTIDENTIFIER']._serialized_end=195
   _globals['_CLUSTERIDENTIFIER']._serialized_start=197
@@ -68,4 +84,10 @@ if _descriptor._USE_C_DESCRIPTORS == False:
   _globals['_MANAGEDCLUSTERIDENTIFIER']._serialized_end=859
   _globals['_POLICYIDENTIFIER']._serialized_start=861
   _globals['_POLICYIDENTIFIER']._serialized_end=944
+  _globals['_RUNIDENTIFIER']._serialized_start=947
+  _globals['_RUNIDENTIFIER']._serialized_end=1094
+  _globals['_ACTIONIDENTIFIER']._serialized_start=1096
+  _globals['_ACTIONIDENTIFIER']._serialized_end=1205
+  _globals['_ACTIONATTEMPTIDENTIFIER']._serialized_start=1208
+  _globals['_ACTIONATTEMPTIDENTIFIER']._serialized_end=1342
 # @@protoc_insertion_point(module_scope)