flyte 0.2.0b1__py3-none-any.whl → 2.0.0b46__py3-none-any.whl

flyte/_internal/controllers/remote/_action.py

@@ -1,11 +1,20 @@
 from __future__ import annotations
 
 from dataclasses import dataclass
+from typing import Literal, Optional
 
-from flyteidl.core import execution_pb2
+from flyteidl2.common import identifier_pb2, phase_pb2
+from flyteidl2.core import execution_pb2, interface_pb2
+from flyteidl2.task import common_pb2, task_definition_pb2
+from flyteidl2.workflow import (
+    run_definition_pb2,
+    state_service_pb2,
+)
+from google.protobuf import timestamp_pb2
 
-from flyte._datastructures import GroupData
-from flyte._protos.workflow import run_definition_pb2, state_service_pb2, task_definition_pb2
+from flyte.models import GroupData
+
+ActionType = Literal["task", "trace"]
 
 
 @dataclass
@@ -15,19 +24,23 @@ 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: run_definition_pb2.TraceAction | None = None
     inputs_uri: str | None = None
     run_output_base: str | None = None
     realized_outputs_uri: str | None = None
     err: execution_pb2.ExecutionError | None = None
-    phase: run_definition_pb2.Phase | None = None
+    phase: phase_pb2.ActionPhase | None = None
     started: bool = False
     retries: int = 0
+    queue: Optional[str] = None  # The queue to which this action was submitted.
     client_err: Exception | None = None  # This error is set when something goes wrong in the controller.
+    cache_key: str | None = None  # None means no caching, otherwise it is the version of the cache.
 
     @property
     def name(self) -> str:
@@ -42,10 +55,10 @@ class Action:
         if self.phase is None:
             return False
         return self.phase in [
-            run_definition_pb2.Phase.PHASE_FAILED,
-            run_definition_pb2.Phase.PHASE_SUCCEEDED,
-            run_definition_pb2.Phase.PHASE_ABORTED,
-            run_definition_pb2.Phase.PHASE_TIMED_OUT,
+            phase_pb2.ACTION_PHASE_FAILED,
+            phase_pb2.ACTION_PHASE_SUCCEEDED,
+            phase_pb2.ACTION_PHASE_ABORTED,
+            phase_pb2.ACTION_PHASE_TIMED_OUT,
         ]
 
     def increment_retries(self):
@@ -61,7 +74,7 @@ class Action:
 
     def mark_cancelled(self):
         self.mark_started()
-        self.phase = run_definition_pb2.Phase.PHASE_ABORTED
+        self.phase = phase_pb2.ACTION_PHASE_ABORTED
 
     def merge_state(self, obj: state_service_pb2.ActionUpdate):
         """
@@ -91,6 +104,8 @@ class Action:
         if not self.started:
             self.task = action.task
 
+        self.cache_key = action.cache_key
+
     def set_client_error(self, exc: Exception):
         self.client_err = exc
 
@@ -101,11 +116,13 @@ 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,
         run_output_base: str,
+        cache_key: str | None = None,
+        queue: Optional[str] = None,
     ) -> Action:
         return cls(
             action_id=sub_action_id,
@@ -115,6 +132,8 @@ class Action:
             task=task_spec,
             inputs_uri=inputs_uri,
             run_output_base=run_output_base,
+            cache_key=cache_key,
+            queue=queue,
         )
 
     @classmethod
@@ -130,7 +149,7 @@ class Action:
         """
         from flyte._logging import logger
 
-        logger.info(f"In Action from_state {obj.action_id} {obj.phase} {obj.output_uri}")
+        logger.debug(f"In Action from_state {obj.action_id} {obj.phase} {obj.output_uri}")
         return cls(
             action_id=obj.action_id,
             parent_action_name=parent_action_name,
@@ -139,3 +158,54 @@ class Action:
             err=obj.error if obj.HasField("error") else None,
             realized_outputs_uri=obj.output_uri,
         )
+
+    @classmethod
+    def from_trace(
+        cls,
+        parent_action_name: str,
+        action_id: identifier_pb2.ActionIdentifier,
+        friendly_name: str,
+        group_data: GroupData | None,
+        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,
+        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.
+        """
+        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)
+
+        spec = task_definition_pb2.TraceSpec(interface=typed_interface) if typed_interface else None
+
+        return cls(
+            action_id=action_id,
+            parent_action_name=parent_action_name,
+            type="trace",
+            friendly_name=friendly_name,
+            group=group_data,
+            inputs_uri=inputs_uri,
+            realized_outputs_uri=outputs_uri,
+            phase=phase_pb2.ACTION_PHASE_SUCCEEDED,
+            run_output_base=run_output_base,
+            trace=run_definition_pb2.TraceAction(
+                name=friendly_name,
+                phase=phase_pb2.ACTION_PHASE_SUCCEEDED,
+                start_time=st,
+                end_time=et,
+                outputs=common_pb2.OutputReferences(
+                    output_uri=outputs_uri,
+                    report_uri=report_uri,
+                ),
+                spec=spec,
+            ),
+        )

flyte/_internal/controllers/remote/_client.py

@@ -1,8 +1,8 @@
 from __future__ import annotations
 
 import grpc.aio
+from flyteidl2.workflow import queue_service_pb2_grpc, state_service_pb2_grpc
 
-from flyte._protos.workflow import queue_service_pb2_grpc, state_service_pb2_grpc
 from flyte.remote import create_channel
 
 from ._service_protocol import QueueService, StateService
@@ -20,7 +20,11 @@ class ControllerClient:
 
     @classmethod
     async def for_endpoint(cls, endpoint: str, insecure: bool = False, **kwargs) -> ControllerClient:
-        return cls(await create_channel(endpoint, insecure=insecure, **kwargs))
+        return cls(await create_channel(endpoint, None, insecure=insecure, **kwargs))
+
+    @classmethod
+    async def for_api_key(cls, api_key: str, insecure: bool = False, **kwargs) -> ControllerClient:
+        return cls(await create_channel(None, api_key, insecure=insecure, **kwargs))
 
     @property
     def state_service(self) -> StateService: