flyte 0.0.1b0__py3-none-any.whl → 2.0.0b46__py3-none-any.whl

flyte/remote/_run.py

@@ -1,92 +1,25 @@
 from __future__ import annotations
 
-import asyncio
 from dataclasses import dataclass, field
-from datetime import datetime, timedelta, timezone
-from typing import AsyncGenerator, AsyncIterator, Iterator, Literal, Tuple, Union, cast
+from typing import AsyncGenerator, AsyncIterator, Literal, Tuple
 
 import grpc
 import rich.repr
-from google.protobuf import timestamp
-from rich.console import Console
-from rich.progress import Progress, SpinnerColumn, TextColumn, TimeElapsedColumn
+from flyteidl2.common import identifier_pb2, list_pb2, phase_pb2
+from flyteidl2.workflow import run_definition_pb2, run_service_pb2
 
-from flyte._api_commons import syncer
-from flyte._initialize import get_client, get_common_config, requires_client
-from flyte._protos.common import identifier_pb2, list_pb2
-from flyte._protos.workflow import run_definition_pb2, run_service_pb2
+from flyte._initialize import ensure_client, get_client, get_init_config
+from flyte._logging import logger
+from flyte.models import ActionPhase
+from flyte.syncify import syncify
 
-from .._protos.workflow.run_service_pb2 import WatchActionDetailsResponse
-from ._console import get_run_url
-from ._logs import Logs
-
-
-def _action_time_phase(action: run_definition_pb2.Action | run_definition_pb2.ActionDetails) -> rich.repr.Result:
-    """
-    Rich representation of the action time and phase.
-    """
-    start_time = timestamp.to_datetime(action.status.start_time, timezone.utc)
-    yield "start_time", start_time.isoformat()
-    if action.status.phase in [
-        run_definition_pb2.PHASE_FAILED,
-        run_definition_pb2.PHASE_SUCCEEDED,
-        run_definition_pb2.PHASE_ABORTED,
-        run_definition_pb2.PHASE_TIMED_OUT,
-    ]:
-        end_time = timestamp.to_datetime(action.status.end_time, timezone.utc)
-        yield "end_time", end_time.isoformat()
-        yield "run_time", f"{(end_time - start_time).seconds} secs"
-    else:
-        yield "end_time", None
-        yield "run_time", f"{(datetime.now(timezone.utc) - start_time).seconds} secs"
-    yield "phase", run_definition_pb2.Phase.Name(action.status.phase)
-    if isinstance(action, run_definition_pb2.ActionDetails):
-        yield "error", action.error_info if action.HasField("error_info") else "NA"
-
-
-def _action_rich_repr(action: run_definition_pb2.Action, root: bool = False) -> rich.repr.Result:
-    """
-    Rich representation of the action.
-    """
-    yield "run-name", action.id.run.name
-    yield "name", action.id.name
-    yield from _action_time_phase(action)
-    yield "task", action.metadata.task.id.name
-    if not root:
-        yield "group", action.metadata.group
-        yield "parent", action.metadata.parent
-
-
-def _action_details_rich_repr(action: run_definition_pb2.ActionDetails, root: bool = False) -> rich.repr.Result:
-    """
-    Rich representation of the action details.
-    """
-    yield "name", action.id.run.name
-    yield from _action_time_phase(action)
-    # yield "task", action.metadata.task.id.name
-    yield "task", action.resolved_task_spec.task_template.id.name
-    yield "task_type", action.resolved_task_spec.task_template.type
-    yield "task_version", action.resolved_task_spec.task_template.id.version
-    if not root:
-        yield "group", action.metadata.group
-        yield "parent", action.metadata.parent
-    # TODO attempt info
-
-
-def _action_done_check(phase: run_definition_pb2.Phase) -> bool:
-    """
-    Check if the action is done.
-    """
-    return phase in [
-        run_definition_pb2.PHASE_FAILED,
-        run_definition_pb2.PHASE_SUCCEEDED,
-        run_definition_pb2.PHASE_ABORTED,
-        run_definition_pb2.PHASE_TIMED_OUT,
-    ]
+from . import Action, ActionDetails, ActionInputs, ActionOutputs
+from ._action import _action_details_rich_repr, _action_rich_repr
+from ._common import ToJSONMixin, filtering, sorting
 
 
 @dataclass
-class Run:
+class Run(ToJSONMixin):
     """
     A class representing a run of a task. It is used to manage the run of a task and its state on the remote
     Union API.
@@ -104,32 +37,84 @@ class Run:
             raise RuntimeError("Run does not have an action")
         self.action = Action(self.pb2.action)
 
+    @syncify
     @classmethod
-    @requires_client
-    @syncer.wrap
     async def listall(
         cls,
-        filters: str | None = None,
+        in_phase: Tuple[ActionPhase | str, ...] | None = None,
+        task_name: str | None = None,
+        task_version: str | None = None,
+        created_by_subject: str | None = None,
         sort_by: Tuple[str, Literal["asc", "desc"]] | None = None,
-    ) -> Union[Iterator[Run], AsyncGenerator[Run, None]]:
+        limit: int = 100,
+    ) -> AsyncIterator[Run]:
         """
         Get all runs for the current project and domain.
 
-        :param filters: The filters to apply to the project list.
-        :param sort_by: The sorting criteria for the project list, in the format (field, order).
+        :param in_phase: Filter runs by one or more phases.
+        :param task_name: Filter runs by task name.
+        :param task_version: Filter runs by task version.
+        :param created_by_subject: Filter runs by the subject that created them. (this is not username, but the subject)
+        :param sort_by: The sorting criteria for the Run list, in the format (field, order).
+        :param limit: The maximum number of runs to return.
         :return: An iterator of runs.
         """
+        ensure_client()
         token = None
-        sort_by = sort_by or ("created_at", "asc")
-        sort_pb2 = list_pb2.Sort(
-            key=sort_by[0], direction=list_pb2.Sort.ASCENDING if sort_by[1] == "asc" else list_pb2.Sort.DESCENDING
-        )
-        cfg = get_common_config()
+        sort_pb2 = sorting(sort_by)
+        filters = []
+        if in_phase:
+            phases = [
+                str(p.to_protobuf_value())
+                if isinstance(p, ActionPhase)
+                else str(phase_pb2.ActionPhase.Value(f"ACTION_PHASE_{p.upper()}"))
+                for p in in_phase
+            ]
+            logger.debug(f"Fetching run phases: {phases}")
+            if len(phases) > 1:
+                filters.append(
+                    list_pb2.Filter(
+                        function=list_pb2.Filter.Function.VALUE_IN,
+                        field="phase",
+                        values=phases,
+                    ),
+                )
+            else:
+                filters.append(
+                    list_pb2.Filter(
+                        function=list_pb2.Filter.Function.EQUAL,
+                        field="phase",
+                        values=phases[0],
+                    ),
+                )
+
+        if task_name:
+            filters.append(
+                list_pb2.Filter(
+                    function=list_pb2.Filter.Function.EQUAL,
+                    field="task_name",
+                    values=[task_name],
+                ),
+            )
+        if task_version:
+            filters.append(
+                list_pb2.Filter(
+                    function=list_pb2.Filter.Function.EQUAL,
+                    field="task_version",
+                    values=[task_version],
+                ),
+            )
+
+        filters = filtering(created_by_subject, *filters)
+
+        cfg = get_init_config()
+        i = 0
         while True:
             req = list_pb2.ListRequest(
-                limit=100,
+                limit=min(100, limit),
                 token=token,
                 sort_by=sort_pb2,
+                filters=filters,
             )
             resp = await get_client().run_service.ListRuns(
                 run_service_pb2.ListRunsRequest(
@@ -144,20 +129,23 @@ class Run:
             )
             token = resp.token
             for r in resp.runs:
+                i += 1
+                if i > limit:
+                    return
                 yield cls(r)
             if not token:
                 break
 
+    @syncify
     @classmethod
-    @requires_client
-    @syncer.wrap
     async def get(cls, name: str) -> Run:
         """
         Get the current run.
 
         :return: The current run.
         """
-        run_details: RunDetails = await RunDetails.get.aio(RunDetails, name=name)
+        ensure_client()
+        run_details: RunDetails = await RunDetails.get.aio(name=name)
         run = run_definition_pb2.Run(
             action=run_definition_pb2.Action(
                 id=run_details.action_id,
@@ -179,112 +167,104 @@ class Run:
         """
         Get the phase of the run.
         """
-        return run_definition_pb2.Phase.Name(self.action.phase)
+        return self.action.phase
+
+    @property
+    def raw_phase(self) -> phase_pb2.ActionPhase:
+        """
+        Get the raw phase of the run.
+        """
+        return self.action.raw_phase
 
-    @syncer.wrap
-    async def wait(self, quiet: bool = False) -> None:
+    @syncify
+    async def wait(self, quiet: bool = False, wait_for: Literal["terminal", "running"] = "terminal") -> None:
         """
         Wait for the run to complete, displaying a rich progress panel with status transitions,
         time elapsed, and error details in case of failure.
-        """
-        console = Console()
-        if self.done():
-            if not quiet:
-                console.print(f"[bold green]Run '{self.name}' is already completed.[/bold green]")
-            return
 
-        try:
-            with Progress(
-                SpinnerColumn(),
-                TextColumn("[progress.description]{task.description}"),
-                TimeElapsedColumn(),
-                console=console,
-                transient=True,
-                disable=quiet,
-            ) as progress:
-                task_id = progress.add_task(f"Waiting for run '{self.name}'...", start=False)
-
-                async for ad in self.watch(cache_data_on_done=True):
-                    if ad is None:
-                        break
-
-                    # Update progress description with the current phase
-                    progress.update(
-                        task_id,
-                        description=f"Run: {self.name} in {ad.phase}, Runtime: {ad.runtime} secs "
-                        f"Attempts[{ad.attempts}]",
-                    )
-                    progress.start_task(task_id)
-
-                    # If the action is done, handle the final state
-                    if ad.done():
-                        progress.stop_task(task_id)
-                        if ad.pb2.status.phase == run_definition_pb2.PHASE_SUCCEEDED:
-                            console.print(f"[bold green]Run '{self.name}' completed successfully.[/bold green]")
-                        else:
-                            console.print(
-                                f"[bold red]Run '{self.name}' exited unsuccessfully in state {ad.phase}"
-                                f"with error: {ad.error_info}[/bold red]"
-                            )
-                        break
-        except asyncio.CancelledError:
-            # Handle cancellation gracefully
-            pass
-        except KeyboardInterrupt:
-            # Handle keyboard interrupt gracefully
-            console.print(f"\n[bold yellow]Run '{self.name}' was interrupted.[/bold yellow]")
+        This method updates the Run's internal state, ensuring that properties like
+        `run.action.phase` reflect the final state after waiting completes.
+        """
+        await self.action.wait(quiet=quiet, wait_for=wait_for)
+        # Update the Run's pb2.action to keep it in sync after waiting
+        self.pb2.action.CopyFrom(self.action.pb2)
 
     async def watch(self, cache_data_on_done: bool = False) -> AsyncGenerator[ActionDetails, None]:
         """
-        Get the details of the run. This is a placeholder for getting the run details.
+        Watch the run for updates, updating the internal Run state with latest details.
+
+        This method updates the Run's action state, ensuring that properties like
+        `run.action.phase` reflect the current state after watching.
         """
-        async for ad in self.action.watch_details(cache_data_on_done=cache_data_on_done):
-            if ad is None:
-                return
+        async for ad in self.action.watch(cache_data_on_done=cache_data_on_done):
+            # The action's pb2 is already updated by Action.watch()
+            # Update the Run's pb2.action to keep it in sync
+            self.pb2.action.CopyFrom(self.action.pb2)
             yield ad
 
-    async def show_logs(self, attempt: int = 1, max_lines: int = 100, show_ts: bool = False, raw: bool = False):
-        return await Logs.create_viewer(
-            action_id=self.action.action_id,
-            attempt=attempt,
-            max_lines=max_lines,
-            show_ts=show_ts,
-            raw=raw,
-        )
+    @syncify
+    async def show_logs(
+        self,
+        attempt: int | None = None,
+        max_lines: int = 100,
+        show_ts: bool = False,
+        raw: bool = False,
+        filter_system: bool = False,
+    ):
+        await self.action.show_logs.aio(attempt, max_lines, show_ts, raw, filter_system=filter_system)
 
+    @syncify
     async def details(self) -> RunDetails:
         """
         Get the details of the run. This is a placeholder for getting the run details.
         """
-        if self._details is None:
-            self._details = await RunDetails.get_details(RunDetails, self.pb2.action.id.run)
+        if self._details is None or not self._details.done():
+            self._details = await RunDetails.get_details.aio(self.pb2.action.id.run)
         return self._details
 
+    @syncify
+    async def inputs(self) -> ActionInputs:
+        """
+        Get the inputs of the run. This is a placeholder for getting the run inputs.
+        """
+        details = await self.details.aio()
+        return await details.inputs()
+
+    @syncify
+    async def outputs(self) -> ActionOutputs:
+        """
+        Get the outputs of the run. This is a placeholder for getting the run outputs.
+        """
+        details = await self.details.aio()
+        return await details.outputs()
+
     @property
-    @requires_client
     def url(self) -> str:
         """
         Get the URL of the run.
         """
         client = get_client()
-        return get_run_url(
-            client.endpoint,
-            insecure=client.insecure,
+        return client.console.run_url(
             project=self.pb2.action.id.run.project,
             domain=self.pb2.action.id.run.domain,
             run_name=self.name,
         )
 
-    @syncer.wrap
-    async def cancel(self) -> None:
+    @syncify
+    async def abort(self):
         """
-        Cancel the run.
+        Aborts / Terminates the run.
         """
-        await get_client().run_service.AbortRun(
-            run_service_pb2.AbortRunRequest(
-                run_id=self.pb2.action.id.run,
+        try:
+            await get_client().run_service.AbortRun(
+                run_service_pb2.AbortRunRequest(
+                    run_id=self.pb2.action.id.run,
+                )
             )
-        )
+        except grpc.aio.AioRpcError as e:
+            if e.code() == grpc.StatusCode.NOT_FOUND:
+                return
+            raise
 
     def done(self) -> bool:
         """
@@ -304,7 +284,8 @@ class Run:
         """
         Rich representation of the Run object.
         """
-        yield from _action_rich_repr(self.pb2.action, root=True)
+        yield "url", f"[blue bold][link={self.url}]link[/link][/blue bold]"
+        yield from _action_rich_repr(self.pb2.action)
 
     def __repr__(self) -> str:
         """
@@ -316,7 +297,7 @@ class Run:
 
 
 @dataclass
-class RunDetails:
+class RunDetails(ToJSONMixin):
     """
     A class representing a run of a task. It is used to manage the run of a task and its state on the remote
     Union API.
@@ -331,13 +312,13 @@ class RunDetails:
         """
         self.action_details = ActionDetails(self.pb2.action)
 
+    @syncify
     @classmethod
-    @requires_client
-    @syncer.wrap
-    async def get_details(cls, run_id: run_definition_pb2.RunIdentifier) -> RunDetails:
+    async def get_details(cls, run_id: identifier_pb2.RunIdentifier) -> RunDetails:
         """
         Get the details of the run. This is a placeholder for getting the run details.
         """
+        ensure_client()
         resp = await get_client().run_service.GetRunDetails(
             run_service_pb2.GetRunDetailsRequest(
                 run_id=run_id,
@@ -345,9 +326,8 @@ class RunDetails:
         )
         return cls(resp.details)
 
+    @syncify
     @classmethod
-    @requires_client
-    @syncer.wrap
     async def get(cls, name: str | None = None) -> RunDetails:
         """
         Get a run by its ID or name. If both are provided, the ID will take precedence.
@@ -355,10 +335,10 @@ class RunDetails:
         :param uri: The URI of the run.
         :param name: The name of the run.
         """
-        cfg = get_common_config()
+        ensure_client()
+        cfg = get_init_config()
         return await RunDetails.get_details.aio(
-            cls,
-            run_id=run_definition_pb2.RunIdentifier(
+            run_id=identifier_pb2.RunIdentifier(
                 org=cfg.org,
                 project=cfg.project,
                 domain=cfg.domain,
@@ -381,7 +361,7 @@ class RunDetails:
         return self.action_details.task_name
 
     @property
-    def action_id(self) -> run_definition_pb2.ActionIdentifier:
+    def action_id(self) -> identifier_pb2.ActionIdentifier:
         """
         Get the action ID.
         """
@@ -410,196 +390,12 @@ class RunDetails:
         """
         Rich representation of the Run object.
         """
-        yield from _action_details_rich_repr(self.pb2.action, root=True)
-
-    def __repr__(self) -> str:
-        """
-        String representation of the Action object.
-        """
-        import rich.pretty
-
-        return rich.pretty.pretty_repr(self)
-
-
-@dataclass
-class Action:
-    """
-    A class representing an action. It is used to manage the run of a task and its state on the remote Union API.
-    """
-
-    pb2: run_definition_pb2.Action
-    _details: ActionDetails | None = None
-
-    @classmethod
-    @requires_client
-    @syncer.wrap
-    async def listall(
-        cls,
-        for_run_name: str,
-        filters: str | None = None,
-        sort_by: Tuple[str, Literal["asc", "desc"]] | None = None,
-    ) -> Union[Iterator[Action], AsyncGenerator[Action, None]]:
-        """
-        Get all actions for a given run.
-
-        :param for_run_name: The name of the run.
-        :param filters: The filters to apply to the project list.
-        :param sort_by: The sorting criteria for the project list, in the format (field, order).
-        :return: An iterator of projects.
-        """
-        token = None
-        sort_by = sort_by or ("created_at", "asc")
-        sort_pb2 = list_pb2.Sort(
-            key=sort_by[0], direction=list_pb2.Sort.ASCENDING if sort_by[1] == "asc" else list_pb2.Sort.DESCENDING
-        )
-        cfg = get_common_config()
-        while True:
-            req = list_pb2.ListRequest(
-                limit=100,
-                token=token,
-                sort_by=sort_pb2,
-            )
-            resp = await get_client().run_service.ListActions(
-                run_service_pb2.ListActionsRequest(
-                    request=req,
-                    run_id=run_definition_pb2.RunIdentifier(
-                        org=cfg.org,
-                        project=cfg.project,
-                        domain=cfg.domain,
-                        name=for_run_name,
-                    ),
-                )
-            )
-            token = resp.token
-            for r in resp.actions:
-                yield cls(r)
-            if not token:
-                break
-
-    @classmethod
-    @requires_client
-    @syncer.wrap
-    async def get(cls, uri: str | None = None, /, run_name: str | None = None, name: str | None = None) -> Action:
-        """
-        Get a run by its ID or name. If both are provided, the ID will take precedence.
-
-        :param uri: The URI of the action.
-        :param run_name: The name of the action.
-        :param name: The name of the action.
-        """
-        cfg = get_common_config()
-        details: ActionDetails = await ActionDetails.get_details.aio(
-            cls,
-            run_definition_pb2.ActionIdentifier(
-                run=run_definition_pb2.RunIdentifier(
-                    org=cfg.org,
-                    project=cfg.project,
-                    domain=cfg.domain,
-                    name=run_name,
-                ),
-                name=name,
-            ),
-        )
-        return cls(
-            pb2=run_definition_pb2.Action(
-                id=details.action_id,
-                metadata=details.pb2.metadata,
-                status=details.pb2.status,
-            ),
-            _details=details,
-        )
-
-    @property
-    def phase(self) -> str:
-        """
-        Get the phase of the action.
-        """
-        return run_definition_pb2.Phase.Name(self.pb2.status.phase)
-
-    @property
-    def name(self) -> str:
-        """
-        Get the name of the action.
-        """
-        return self.action_id.name
-
-    @property
-    def run_name(self) -> str:
-        """
-        Get the name of the run.
-        """
-        return self.action_id.run.name
-
-    @property
-    def task_name(self) -> str | None:
-        """
-        Get the name of the task.
-        """
-        if self.pb2.metadata.HasField("task") and self.pb2.metadata.task.HasField("id"):
-            return self.pb2.metadata.task.id.name
-        return None
-
-    @property
-    def action_id(self) -> run_definition_pb2.ActionIdentifier:
-        """
-        Get the action ID.
-        """
-        return self.pb2.id
-
-    async def show_logs(
-        self, attempt: int | None = None, max_lines: int = 30, show_ts: bool = False, raw: bool = False
-    ):
-        details = await self.details()
-        if not attempt:
-            attempt = details.attempts
-        if details.phase in [
-            run_definition_pb2.PHASE_QUEUED,
-            run_definition_pb2.PHASE_INITIALIZING,
-            run_definition_pb2.PHASE_WAITING_FOR_RESOURCES,
-        ]:
-            raise RuntimeError("Action has not yet started, so logs are not available.")
-        return await Logs.create_viewer(
-            action_id=self.action_id, attempt=attempt, max_lines=max_lines, show_ts=show_ts, raw=raw
-        )
-
-    async def details(self) -> ActionDetails:
-        """
-        Get the details of the action. This is a placeholder for getting the action details.
-        """
-        if not self._details:
-            self._details = await ActionDetails.get_details.aio(ActionDetails, self.action_id)
-        return cast(ActionDetails, self._details)
-
-    async def watch_details(self, cache_data_on_done: bool = False) -> AsyncGenerator[ActionDetails, None]:
-        """
-        Watch the action for updates. This is a placeholder for watching the action.
-        """
-        ad = None
-        async for ad in ActionDetails.watch.aio(ActionDetails, self.action_id):
-            if ad is None:
-                return
-            self._details = ad
-            yield ad
-        if cache_data_on_done and ad and ad.done():
-            await cast(ActionDetails, self._details).outputs()
-
-    def done(self) -> bool:
-        """
-        Check if the action is done.
-        """
-        return _action_done_check(self.pb2.status.phase)
-
-    async def sync(self) -> Action:
-        """
-        Sync the action with the remote server. This is a placeholder for syncing the action.
-        """
-        return self
-
-    def __rich_repr__(self) -> rich.repr.Result:
-        """
-        Rich representation of the Action object.
-        """
-        yield from _action_rich_repr(self.pb2, root=True)
+        yield "labels", str(self.pb2.run_spec.labels)
+        yield "annotations", str(self.pb2.run_spec.annotations)
+        yield "env-vars", str(self.pb2.run_spec.envs)
+        yield "is-interruptible", str(self.pb2.run_spec.interruptible)
+        yield "cache-overwrite", self.pb2.run_spec.overwrite_cache
+        yield from _action_details_rich_repr(self.pb2.action)
 
     def __repr__(self) -> str:
         """
@@ -608,266 +404,3 @@ class Action:
         import rich.pretty
 
         return rich.pretty.pretty_repr(self)
-
-
-@dataclass
-class ActionDetails:
-    """
-    A class representing an action. It is used to manage the run of a task and its state on the remote Union API.
-    """
-
-    pb2: run_definition_pb2.ActionDetails
-    _inputs: ActionInputs | None = None
-    _outputs: ActionOutputs | None = None
-
-    @classmethod
-    @requires_client
-    @syncer.wrap
-    async def get_details(cls, action_id: run_definition_pb2.ActionIdentifier) -> ActionDetails:
-        """
-        Get the details of the action. This is a placeholder for getting the action details.
-        """
-        resp = await get_client().run_service.GetActionDetails(
-            run_service_pb2.GetActionDetailsRequest(
-                action_id=action_id,
-            )
-        )
-        return ActionDetails(resp.details)
-
-    @classmethod
-    @requires_client
-    @syncer.wrap
-    async def get(
-        cls, uri: str | None = None, /, run_name: str | None = None, name: str | None = None
-    ) -> ActionDetails:
-        """
-        Get a run by its ID or name. If both are provided, the ID will take precedence.
-
-        :param uri: The URI of the action.
-        :param name: The name of the action.
-        :param run_name: The name of the run.
-        """
-        if not uri:
-            assert name is not None and run_name is not None, "Either uri or name and run_name must be provided"
-        cfg = get_common_config()
-        return await cls.get_details.aio(
-            cls,
-            run_definition_pb2.ActionIdentifier(
-                run=run_definition_pb2.RunIdentifier(
-                    org=cfg.org,
-                    project=cfg.project,
-                    domain=cfg.domain,
-                    name=run_name,
-                ),
-                name=name,
-            ),
-        )
-
-    @classmethod
-    @requires_client
-    @syncer.wrap
-    async def watch(cls, action_id: run_definition_pb2.ActionIdentifier) -> AsyncGenerator[ActionDetails, None]:
-        """
-        Watch the action for updates. This is a placeholder for watching the action.
-        """
-        if not action_id:
-            raise ValueError("Action ID is required")
-
-        call = cast(
-            AsyncIterator[WatchActionDetailsResponse],
-            get_client().run_service.WatchActionDetails(
-                request=run_service_pb2.WatchActionDetailsRequest(
-                    action_id=action_id,
-                )
-            ),
-        )
-        try:
-            async for resp in call:
-                v = cls(resp.details)
-                yield v
-                if v.done():
-                    return
-        except grpc.aio.AioRpcError as e:
-            if e.code() == grpc.StatusCode.CANCELLED:
-                pass
-            else:
-                raise e
-
-    async def watch_updates(self, cache_data_on_done: bool = False) -> AsyncGenerator[ActionDetails, None]:
-        async for d in self.watch.aio(action_id=self.pb2.id):
-            yield d
-            if d.done():
-                self.pb2 = d.pb2
-                break
-
-        if cache_data_on_done and self.done():
-            await self._cache_data.aio(self)
-
-    @property
-    def phase(self) -> str:
-        """
-        Get the phase of the action.
-        """
-        return run_definition_pb2.Phase.Name(self.status.phase)
-
-    @property
-    def name(self) -> str:
-        """
-        Get the name of the action.
-        """
-        return self.action_id.name
-
-    @property
-    def run_name(self) -> str:
-        """
-        Get the name of the run.
-        """
-        return self.action_id.run.name
-
-    @property
-    def task_name(self) -> str | None:
-        """
-        Get the name of the task.
-        """
-        if self.pb2.metadata.HasField("task") and self.pb2.metadata.task.HasField("id"):
-            return self.pb2.metadata.task.id.name
-        return None
-
-    @property
-    def action_id(self) -> run_definition_pb2.ActionIdentifier:
-        """
-        Get the action ID.
-        """
-        return self.pb2.id
-
-    @property
-    def metadata(self) -> run_definition_pb2.ActionMetadata:
-        return self.pb2.metadata
-
-    @property
-    def status(self) -> run_definition_pb2.ActionStatus:
-        return self.pb2.status
-
-    @property
-    def error_info(self) -> run_definition_pb2.ErrorInfo | None:
-        if self.pb2.HasField("error_info"):
-            return self.pb2.error_info
-        return None
-
-    @property
-    def abort_info(self) -> run_definition_pb2.AbortInfo | None:
-        if self.pb2.HasField("abort_info"):
-            return self.pb2.abort_info
-        return None
-
-    @property
-    def runtime(self) -> timedelta:
-        """
-        Get the runtime of the action.
-        """
-        start_time = timestamp.to_datetime(self.pb2.status.start_time, timezone.utc)
-        if self.pb2.status.HasField("end_time"):
-            end_time = timestamp.to_datetime(self.pb2.status.end_time, timezone.utc)
-            return end_time - start_time
-        return datetime.now(timezone.utc) - start_time
-
-    @property
-    def attempts(self) -> int:
-        """
-        Get the number of attempts of the action.
-        """
-        return self.pb2.status.attempts
-
-    @syncer.wrap
-    async def _cache_data(self) -> bool:
-        """
-        Cache the inputs and outputs of the action.
-        :return: Returns True if Action is terminal and all data is cached else False.
-        """
-        if self._inputs and self._outputs:
-            return True
-        if self._inputs and not self.done():
-            return False
-        resp = await get_client().run_service.GetActionData(
-            request=run_service_pb2.GetActionDataRequest(
-                action_id=self.pb2.id,
-            )
-        )
-        self._inputs = ActionInputs(resp.inputs)
-        self._outputs = ActionOutputs(resp.outputs) if resp.HasField("outputs") else None
-        return self._outputs is not None
-
-    async def inputs(self) -> ActionInputs:
-        """
-        Placeholder for inputs. This can be extended to handle inputs from the run context.
-        """
-        if not self._inputs:
-            await self._cache_data.aio(self)
-        return cast(ActionInputs, self._inputs)
-
-    async def outputs(self) -> ActionOutputs:
-        """
-        Placeholder for outputs. This can be extended to handle outputs from the run context.
-        """
-        if not self._outputs:
-            if not await self._cache_data.aio(self):
-                raise RuntimeError(
-                    "Action is not in a terminal state, outputs are not available. "
-                    "Please wait for the action to complete."
-                )
-        return cast(ActionOutputs, self._outputs)
-
-    def done(self) -> bool:
-        """
-        Check if the action is in a terminal state (completed or failed). This is a placeholder for checking the
-        action state.
-        """
-        return _action_done_check(self.pb2.status.phase)
-
-    def __rich_repr__(self) -> rich.repr.Result:
-        """
-        Rich representation of the Action object.
-        """
-        yield from _action_details_rich_repr(self.pb2, root=True)
-
-    def __repr__(self) -> str:
-        """
-        String representation of the Action object.
-        """
-        import rich.pretty
-
-        return rich.pretty.pretty_repr(self)
-
-
-@dataclass
-class ActionInputs:
-    """
-    A class representing the inputs of an action. It is used to manage the inputs of a task and its state on the
-    remote Union API.
-    """
-
-    pb2: run_definition_pb2.Inputs
-
-    def __repr__(self):
-        import rich.pretty
-
-        import flyte.types as types
-
-        return rich.pretty.pretty_repr(types.literal_string_repr(self.pb2))
-
-
-@dataclass
-class ActionOutputs:
-    """
-    A class representing the outputs of an action. It is used to manage the outputs of a task and its state on the
-    remote Union API.
-    """
-
-    pb2: run_definition_pb2.Outputs
-
-    def __repr__(self):
-        import rich.pretty
-
-        import flyte.types as types
-
-        return rich.pretty.pretty_repr(types.literal_string_repr(self.pb2))