flyte 0.0.1b3__py3-none-any.whl → 0.2.0a0__py3-none-any.whl

flyte/remote/_run.py

@@ -3,7 +3,7 @@ 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, Iterator, List, Literal, Tuple, Union, cast
 
 import grpc
 import rich.repr
@@ -11,15 +11,17 @@ from google.protobuf import timestamp
 from rich.console import Console
 from rich.progress import Progress, SpinnerColumn, TextColumn, TimeElapsedColumn
 
-from flyte._api_commons import syncer
-from flyte._initialize import get_client, get_common_config, requires_client
+from flyte._initialize import ensure_client, get_client, get_common_config
 from flyte._protos.common import identifier_pb2, list_pb2
 from flyte._protos.workflow import run_definition_pb2, run_service_pb2
+from flyte.syncify import syncify
 
 from .._protos.workflow.run_service_pb2 import WatchActionDetailsResponse
 from ._console import get_run_url
 from ._logs import Logs
 
+WaitFor = Literal["terminal", "running", "logs-ready"]
+
 
 def _action_time_phase(action: run_definition_pb2.Action | run_definition_pb2.ActionDetails) -> rich.repr.Result:
     """
@@ -41,36 +43,48 @@ def _action_time_phase(action: run_definition_pb2.Action | run_definition_pb2.Ac
         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"
+        yield (
+            "error",
+            f"{action.error_info.kind}: {action.error_info.message}" if action.HasField("error_info") else "NA",
+        )
 
 
-def _action_rich_repr(action: run_definition_pb2.Action, root: bool = False) -> rich.repr.Result:
+def _action_rich_repr(action: run_definition_pb2.Action) -> rich.repr.Result:
     """
     Rich representation of the action.
     """
-    yield "run-name", action.id.run.name
+    yield "run", action.id.run.name
+    if action.metadata.HasField("task"):
+        yield "task", action.metadata.task.id.name
+        yield "type", "task"
     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
+    yield "group", action.metadata.group
+    yield "parent", action.metadata.parent
+    yield "attempts", action.status.attempts
+
 
+def _attempt_rich_repr(action: List[run_definition_pb2.ActionAttempt]) -> rich.repr.Result:
+    for attempt in action:
+        yield "attempt", attempt.attempt
+        yield "phase", run_definition_pb2.Phase.Name(attempt.phase)
+        yield "logs_available", attempt.logs_available
 
-def _action_details_rich_repr(action: run_definition_pb2.ActionDetails, root: bool = False) -> rich.repr.Result:
+
+def _action_details_rich_repr(action: run_definition_pb2.ActionDetails) -> 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
+    yield "attempts", action.attempts
+    yield "error", f"{action.error_info.kind}: {action.error_info.message}" if action.HasField("error_info") else "NA"
+    yield "phase", run_definition_pb2.Phase.Name(action.status.phase)
+    yield "group", action.metadata.group
+    yield "parent", action.metadata.parent
 
 
 def _action_done_check(phase: run_definition_pb2.Phase) -> bool:
@@ -104,14 +118,13 @@ 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,
         sort_by: Tuple[str, Literal["asc", "desc"]] | None = None,
-    ) -> Union[Iterator[Run], AsyncGenerator[Run, None]]:
+    ) -> AsyncIterator[Run]:
         """
         Get all runs for the current project and domain.
 
@@ -119,6 +132,7 @@ class Run:
         :param sort_by: The sorting criteria for the project list, in the format (field, order).
         :return: An iterator of runs.
         """
+        ensure_client()
         token = None
         sort_by = sort_by or ("created_at", "asc")
         sort_pb2 = list_pb2.Sort(
@@ -148,16 +162,16 @@ class Run:
             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,78 +193,38 @@ 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) -> run_definition_pb2.Phase:
+        """
+        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]")
+        return await self.action.wait(quiet=quiet, wait_for=wait_for)
 
     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.
         """
-        async for ad in self.action.watch_details(cache_data_on_done=cache_data_on_done):
-            if ad is None:
-                return
-            yield ad
+        return self.action.watch(cache_data_on_done=cache_data_on_done)
 
-    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,
-        )
+    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(attempt, max_lines, show_ts, raw, filter_system=filter_system)
 
     async def details(self) -> RunDetails:
         """
@@ -261,7 +235,6 @@ class Run:
         return self._details
 
     @property
-    @requires_client
     def url(self) -> str:
         """
         Get the URL of the run.
@@ -275,16 +248,21 @@ class Run:
             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 +282,7 @@ class Run:
         """
         Rich representation of the Run object.
         """
-        yield from _action_rich_repr(self.pb2.action, root=True)
+        yield from _action_rich_repr(self.pb2.action)
 
     def __repr__(self) -> str:
         """
@@ -331,13 +309,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:
         """
         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 +323,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,9 +332,9 @@ class RunDetails:
         :param uri: The URI of the run.
         :param name: The name of the run.
         """
+        ensure_client()
         cfg = get_common_config()
         return await RunDetails.get_details.aio(
-            cls,
             run_id=run_definition_pb2.RunIdentifier(
                 org=cfg.org,
                 project=cfg.project,
@@ -410,7 +387,7 @@ class RunDetails:
         """
         Rich representation of the Run object.
         """
-        yield from _action_details_rich_repr(self.pb2.action, root=True)
+        yield from _action_details_rich_repr(self.pb2.action)
 
     def __repr__(self) -> str:
         """
@@ -430,15 +407,14 @@ class Action:
     pb2: run_definition_pb2.Action
     _details: ActionDetails | None = None
 
+    @syncify
     @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]]:
+    ) -> Union[Iterator[Action], AsyncIterator[Action]]:
         """
         Get all actions for a given run.
 
@@ -447,6 +423,7 @@ class Action:
         :param sort_by: The sorting criteria for the project list, in the format (field, order).
         :return: An iterator of projects.
         """
+        ensure_client()
         token = None
         sort_by = sort_by or ("created_at", "asc")
         sort_pb2 = list_pb2.Sort(
@@ -476,9 +453,8 @@ class Action:
             if not token:
                 break
 
+    @syncify
     @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.
@@ -487,9 +463,9 @@ class Action:
         :param run_name: The name of the action.
         :param name: The name of the action.
         """
+        ensure_client()
         cfg = get_common_config()
         details: ActionDetails = await ActionDetails.get_details.aio(
-            cls,
             run_definition_pb2.ActionIdentifier(
                 run=run_definition_pb2.RunIdentifier(
                     org=cfg.org,
@@ -516,6 +492,13 @@ class Action:
         """
         return run_definition_pb2.Phase.Name(self.pb2.status.phase)
 
+    @property
+    def raw_phase(self) -> run_definition_pb2.Phase:
+        """
+        Get the raw phase of the action.
+        """
+        return self.pb2.status.phase
+
     @property
     def name(self) -> str:
         """
@@ -547,19 +530,27 @@ class Action:
         return self.pb2.id
 
     async def show_logs(
-        self, attempt: int | None = None, max_lines: int = 30, show_ts: bool = False, raw: bool = False
+        self,
+        attempt: int | None = None,
+        max_lines: int = 30,
+        show_ts: bool = False,
+        raw: bool = False,
+        filter_system: bool = False,
     ):
         details = await self.details()
+        if not details.is_running and not details.done():
+            # TODO we can short circuit here if the attempt is not the last one and it is done!
+            await self.wait(wait_for="logs-ready")
+            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
+            action_id=self.action_id,
+            attempt=attempt,
+            max_lines=max_lines,
+            show_ts=show_ts,
+            raw=raw,
+            filter_system=filter_system,
         )
 
     async def details(self) -> ActionDetails:
@@ -567,27 +558,106 @@ class Action:
         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)
+            self._details = await ActionDetails.get_details.aio(self.action_id)
         return cast(ActionDetails, self._details)
 
-    async def watch_details(self, cache_data_on_done: bool = False) -> AsyncGenerator[ActionDetails, None]:
+    async def watch(
+        self, cache_data_on_done: bool = False, wait_for: WaitFor = "terminal"
+    ) -> 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):
+        async for ad in ActionDetails.watch.aio(self.action_id):
             if ad is None:
                 return
             self._details = ad
             yield ad
+            if wait_for == "running" and ad.is_running:
+                break
+            elif wait_for == "logs-ready" and ad.logs_available():
+                break
+            if ad.done():
+                break
         if cache_data_on_done and ad and ad.done():
             await cast(ActionDetails, self._details).outputs()
 
+    async def wait(self, quiet: bool = False, wait_for: WaitFor = "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:
+                if self.pb2.status.phase == run_definition_pb2.PHASE_SUCCEEDED:
+                    console.print(
+                        f"[bold green]Action '{self.name}' in Run '{self.run_name}'"
+                        f" completed successfully.[/bold green]"
+                    )
+                else:
+                    details = await self.details()
+                    console.print(
+                        f"[bold red]Action '{self.name}' in Run '{self.run_name}'"
+                        f" exited unsuccessfully in state {self.phase} with error: {details.error_info}[/bold red]"
+                    )
+            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)
+                progress.start_task(task_id)
+
+                async for ad in self.watch(cache_data_on_done=True, wait_for=wait_for):
+                    if ad is None:
+                        progress.stop_task(task_id)
+                        break
+
+                    if ad.is_running and wait_for == "running":
+                        progress.start_task(task_id)
+                        break
+
+                    if ad.logs_available() and wait_for == "logs-ready":
+                        progress.start_task(task_id)
+                        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}]",
+                    )
+
+                    # 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
+            pass
+
     def done(self) -> bool:
         """
         Check if the action is done.
         """
-        return _action_done_check(self.pb2.status.phase)
+        return _action_done_check(self.raw_phase)
 
     async def sync(self) -> Action:
         """
@@ -599,7 +669,9 @@ class Action:
         """
         Rich representation of the Action object.
         """
-        yield from _action_rich_repr(self.pb2, root=True)
+        yield from _action_rich_repr(self.pb2)
+        if self._details:
+            yield from self._details.__rich_repr__()
 
     def __repr__(self) -> str:
         """
@@ -620,13 +692,13 @@ class ActionDetails:
     _inputs: ActionInputs | None = None
     _outputs: ActionOutputs | None = None
 
+    @syncify
     @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.
         """
+        ensure_client()
         resp = await get_client().run_service.GetActionDetails(
             run_service_pb2.GetActionDetailsRequest(
                 action_id=action_id,
@@ -634,9 +706,8 @@ class ActionDetails:
         )
         return ActionDetails(resp.details)
 
+    @syncify
     @classmethod
-    @requires_client
-    @syncer.wrap
     async def get(
         cls, uri: str | None = None, /, run_name: str | None = None, name: str | None = None
     ) -> ActionDetails:
@@ -647,11 +718,11 @@ class ActionDetails:
         :param name: The name of the action.
         :param run_name: The name of the run.
         """
+        ensure_client()
         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,
@@ -663,13 +734,13 @@ class ActionDetails:
             ),
         )
 
+    @syncify
     @classmethod
-    @requires_client
-    @syncer.wrap
-    async def watch(cls, action_id: run_definition_pb2.ActionIdentifier) -> AsyncGenerator[ActionDetails, None]:
+    async def watch(cls, action_id: run_definition_pb2.ActionIdentifier) -> AsyncIterator[ActionDetails]:
         """
         Watch the action for updates. This is a placeholder for watching the action.
         """
+        ensure_client()
         if not action_id:
             raise ValueError("Action ID is required")
 
@@ -701,7 +772,7 @@ class ActionDetails:
                 break
 
         if cache_data_on_done and self.done():
-            await self._cache_data.aio(self)
+            await self._cache_data.aio()
 
     @property
     def phase(self) -> str:
@@ -710,6 +781,20 @@ class ActionDetails:
         """
         return run_definition_pb2.Phase.Name(self.status.phase)
 
+    @property
+    def raw_phase(self) -> run_definition_pb2.Phase:
+        """
+        Get the raw phase of the action.
+        """
+        return self.status.phase
+
+    @property
+    def is_running(self) -> bool:
+        """
+        Check if the action is currently running.
+        """
+        return self.status.phase == run_definition_pb2.PHASE_RUNNING
+
     @property
     def name(self) -> str:
         """
@@ -778,7 +863,19 @@ class ActionDetails:
         """
         return self.pb2.status.attempts
 
-    @syncer.wrap
+    def logs_available(self, attempt: int | None = None) -> bool:
+        """
+        Check if logs are available for the action, optionally for a specific attempt.
+        If attempt is None, it checks for the latest attempt.
+        """
+        if attempt is None:
+            attempt = self.pb2.status.attempts
+        attempts = self.pb2.attempts
+        if attempts and len(attempts) >= attempt:
+            return attempts[attempt - 1].logs_available
+        return False
+
+    @syncify
     async def _cache_data(self) -> bool:
         """
         Cache the inputs and outputs of the action.
@@ -802,7 +899,7 @@ class ActionDetails:
         Placeholder for inputs. This can be extended to handle inputs from the run context.
         """
         if not self._inputs:
-            await self._cache_data.aio(self)
+            await self._cache_data.aio()
         return cast(ActionInputs, self._inputs)
 
     async def outputs(self) -> ActionOutputs:
@@ -810,7 +907,7 @@ class ActionDetails:
         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):
+            if not await self._cache_data.aio():
                 raise RuntimeError(
                     "Action is not in a terminal state, outputs are not available. "
                     "Please wait for the action to complete."
@@ -822,13 +919,13 @@ class ActionDetails:
         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)
+        return _action_done_check(self.raw_phase)
 
     def __rich_repr__(self) -> rich.repr.Result:
         """
         Rich representation of the Action object.
         """
-        yield from _action_details_rich_repr(self.pb2, root=True)
+        yield from _action_details_rich_repr(self.pb2)
 
     def __repr__(self) -> str:
         """