flyte 0.2.0b2__py3-none-any.whl → 0.2.0b4__py3-none-any.whl

flyte/{_cli → cli}/_run.py

@@ -1,5 +1,6 @@
 from __future__ import annotations
 
+import asyncio
 import inspect
 from dataclasses import dataclass, field, fields
 from pathlib import Path
@@ -50,6 +51,28 @@ class RunArguments:
             )
         },
     )
+    name: str | None = field(
+        default=None,
+        metadata={
+            "click.option": click.Option(
+                ["--name"],
+                type=str,
+                help="Name of the run. If not provided, a random name will be generated.",
+            )
+        },
+    )
+    follow: bool = field(
+        default=True,
+        metadata={
+            "click.option": click.Option(
+                ["--follow", "-f"],
+                is_flag=True,
+                default=False,
+                help="Wait and watch logs for the parent action. If not provided, the cli will exit after "
+                "successfully launching a remote execution with a link to the UI.",
+            )
+        },
+    )
 
     @classmethod
     def from_dict(cls, d: Dict[str, Any]) -> RunArguments:
@@ -76,21 +99,31 @@ class RunTaskCommand(click.Command):
         assert obj.endpoint, "CLI Config should have an endpoint"
         obj.init(self.run_args.project, self.run_args.domain)
 
-        r = flyte.with_runcontext(
-            copy_style=self.run_args.copy_style,
-            version=self.run_args.copy_style,
-            mode="local" if self.run_args.local else "remote",
-        ).run(self.obj, **ctx.params)
-        if isinstance(r, Run) and r.action is not None:
-            console = Console()
-            console.print(
-                common.get_panel(
-                    "Run",
-                    f"[green bold]Created Run: {r.name} [/green bold] "
-                    f"(Project: {r.action.action_id.run.project}, Domain: {r.action.action_id.run.domain})\n\n"
-                    f"[blue bold]{r.url}[/blue bold]",
+        async def _run():
+            r = flyte.with_runcontext(
+                copy_style=self.run_args.copy_style,
+                version=self.run_args.copy_style,
+                mode="local" if self.run_args.local else "remote",
+                name=self.run_args.name,
+            ).run(self.obj, **ctx.params)
+            if isinstance(r, Run) and r.action is not None:
+                console = Console()
+                console.print(
+                    common.get_panel(
+                        "Run",
+                        f"[green bold]Created Run: {r.name} [/green bold] "
+                        f"(Project: {r.action.action_id.run.project}, Domain: {r.action.action_id.run.domain})\n\n"
+                        f"[blue bold]{r.url}[/blue bold]",
+                    )
                 )
-            )
+                if self.run_args.follow:
+                    console.print(
+                        "[dim]Log streaming enabled, will wait for task to start running "
+                        "and log stream to be available[/dim]"
+                    )
+                    await r.show_logs(max_lines=30, show_ts=True, raw=False)
+
+        asyncio.run(_run())
 
     def get_params(self, ctx: Context) -> List[Parameter]:
         # Note this function may be called multiple times by click.
@@ -165,11 +198,11 @@ class TaskFiles(common.FileGroup):
             filename=Path(filename),
             run_args=run_args,
             name=filename,
-            help=f"Run, functions decorated `env.task` or instances of Tasks in {filename}",
+            help=f"Run, functions decorated `env.task` {filename}",
         )
 
 
 run = TaskFiles(
     name="run",
-    help="Run a task from a python file",
+    help="Run a task from a python file.",
 )

flyte/{_cli → cli}/main.py

@@ -3,6 +3,7 @@ import rich_click as click
 from flyte._logging import initialize_logger, logger
 
 from ..config import Config
+from ._abort import abort
 from ._common import CLIConfig
 from ._create import create
 from ._deploy import deploy
@@ -79,7 +80,14 @@ def main(
     config_file: str | None,
 ):
     """
-    v2 cli. Root command, please use one of the subcommands.
+
+ ____  __    _  _  ____  ____    _  _  ____     __
+(  __)(  )  ( \\/ )(_  _)(  __)  / )( \\(___ \\   /  \
+ ) _) / (_/\\ )  /   )(   ) _)   \\ \\/ / / __/ _(  0 )
+(__)  \\____/(__/   (__) (____)   \\__/ (____)(_)\\__/
+
+    The flyte cli follows a simple verb based structure, where the top-level commands are verbs that describe the action
+to be taken, and the subcommands are nouns that describe the object of the action.
     """
     log_level = _verbosity_to_loglevel(verbose)
     if log_level is not None:
@@ -102,3 +110,4 @@ main.add_command(run)
 main.add_command(deploy)
 main.add_command(get)  # type: ignore
 main.add_command(create)  # type: ignore
+main.add_command(abort)  # type: ignore

flyte/config/_config.py

@@ -183,6 +183,8 @@ def get_config_file(c: typing.Union[str, ConfigFile, None]) -> ConfigFile | None
     if isinstance(c, str):
         logger.debug(f"Using specified config file at {c}")
         return ConfigFile(c)
+    elif isinstance(c, ConfigFile):
+        return c
     config_path = resolve_config_path()
     if config_path:
         return ConfigFile(str(config_path))

flyte/errors.py

@@ -141,3 +141,12 @@ class ReferenceTaskError(RuntimeUserError):
 
     def __init__(self, message: str):
         super().__init__("ReferenceTaskUsageError", message, "user")
+
+
+class LogsNotYetAvailableError(BaseRuntimeError):
+    """
+    This error is raised when the logs are not yet available for a task.
+    """
+
+    def __init__(self, message: str):
+        super().__init__("LogsNotYetAvailable", "system", message, None)

flyte/remote/_client/auth/_channel.py

@@ -116,6 +116,7 @@ async def create_channel(
         kwargs["auth_type"] = "ClientSecret"
         kwargs["client_id"] = client_id
         kwargs["client_secret"] = client_secret
+        kwargs["client_credentials_secret"] = client_secret
 
     assert endpoint, "Endpoint must be specified by this point"
 

flyte/remote/_logs.py

@@ -3,6 +3,7 @@ from collections import deque
 from dataclasses import dataclass
 from typing import AsyncGenerator, AsyncIterator
 
+import grpc
 from rich.console import Console
 from rich.live import Live
 from rich.panel import Panel
@@ -12,16 +13,23 @@ from flyte._api_commons import syncer
 from flyte._initialize import get_client, requires_client
 from flyte._protos.logs.dataplane import payload_pb2
 from flyte._protos.workflow import run_definition_pb2, run_logs_service_pb2
+from flyte.errors import LogsNotYetAvailableError
 
+style_map = {
+    payload_pb2.LogLineOriginator.SYSTEM: "bold magenta",
+    payload_pb2.LogLineOriginator.USER: "cyan",
+    payload_pb2.LogLineOriginator.UNKNOWN: "light red",
+}
 
-def _format_line(logline: payload_pb2.LogLine, show_ts: bool) -> Text:
-    style_map = {
-        payload_pb2.LogLineOriginator.SYSTEM: "bold magenta",
-        payload_pb2.LogLineOriginator.USER: "cyan",
-        payload_pb2.LogLineOriginator.UNKNOWN: "light red",
-    }
+
+def _format_line(logline: payload_pb2.LogLine, show_ts: bool, filter_system: bool) -> Text | None:
+    if filter_system:
+        if logline.originator == payload_pb2.LogLineOriginator.SYSTEM:
+            return None
     style = style_map.get(logline.originator, "")
     if "flyte" in logline.message and "flyte.errors" not in logline.message:
+        if filter_system:
+            return None
         style = "dim"
     ts = ""
     if show_ts:
@@ -34,7 +42,14 @@ class AsyncLogViewer:
     A class to view logs asynchronously in the console or terminal or jupyter notebook.
     """
 
-    def __init__(self, log_source: AsyncIterator, max_lines: int = 30, name: str = "Logs", show_ts: bool = False):
+    def __init__(
+        self,
+        log_source: AsyncIterator,
+        max_lines: int = 30,
+        name: str = "Logs",
+        show_ts: bool = False,
+        filter_system: bool = False,
+    ):
         self.console = Console()
         self.log_source = log_source
         self.max_lines = max_lines
@@ -42,23 +57,30 @@ class AsyncLogViewer:
         self.name = name
         self.show_ts = show_ts
         self.total_lines = 0
+        self.filter_flyte = filter_system
 
-    def _render(self):
+    def _render(self) -> Panel:
         log_text = Text()
         for line in self.lines:
             log_text.append(line)
         return Panel(log_text, title=self.name, border_style="yellow")
 
     async def run(self):
-        with Live(self._render(), refresh_per_second=10, console=self.console) as live:
+        with Live(self._render(), refresh_per_second=20, console=self.console) as live:
             try:
                 async for logline in self.log_source:
-                    formatted = _format_line(logline, show_ts=self.show_ts)
-                    self.lines.append(formatted)
+                    formatted = _format_line(logline, show_ts=self.show_ts, filter_system=self.filter_flyte)
+                    if formatted:
+                        self.lines.append(formatted)
                     self.total_lines += 1
                     live.update(self._render())
             except asyncio.CancelledError:
                 pass
+            except KeyboardInterrupt:
+                pass
+            except LogsNotYetAvailableError as e:
+                self.console.print(f"[red]Error:[/red] {e}")
+                live.update("")
         self.console.print(f"Scrolled {self.total_lines} lines of logs.")
 
 
@@ -75,14 +97,24 @@ class Logs:
         :param action_id: The action ID to tail logs for.
         :param attempt: The attempt number (default is 0).
         """
-        resp = get_client().logs_service.TailLogs(
-            run_logs_service_pb2.TailLogsRequest(action_id=action_id, attempt=attempt)
-        )
-        async for log_set in resp:
-            if log_set.logs:
-                for log in log_set.logs:
-                    for line in log.lines:
-                        yield line
+        try:
+            resp = get_client().logs_service.TailLogs(
+                run_logs_service_pb2.TailLogsRequest(action_id=action_id, attempt=attempt)
+            )
+            async for log_set in resp:
+                if log_set.logs:
+                    for log in log_set.logs:
+                        for line in log.lines:
+                            yield line
+        except asyncio.CancelledError:
+            pass
+        except KeyboardInterrupt:
+            pass
+        except grpc.aio.AioRpcError as e:
+            if e.code() == grpc.StatusCode.NOT_FOUND:
+                raise LogsNotYetAvailableError(
+                    f"Log stream not available for action {action_id.name} in run {action_id.run.name}."
+                )
 
     @classmethod
     async def create_viewer(
@@ -92,6 +124,7 @@ class Logs:
         max_lines: int = 30,
         show_ts: bool = False,
         raw: bool = False,
+        filter_system: bool = False,
     ):
         """
         Create a log viewer for a given action ID and attempt.
@@ -101,16 +134,20 @@ class Logs:
            and keep only max_lines in view.
         :param show_ts: Whether to show timestamps in the logs.
         :param raw: if True, return the raw log lines instead of a viewer.
+        :param filter_system: Whether to filter log lines based on system logs.
         """
         if raw:
             console = Console()
             async for line in cls.tail.aio(cls, action_id=action_id, attempt=attempt):
-                console.print(_format_line(line, show_ts=show_ts), end="")
+                line_text = _format_line(line, show_ts=show_ts, filter_system=filter_system)
+                if line_text:
+                    console.print(line_text, end="")
             return
         viewer = AsyncLogViewer(
             log_source=cls.tail.aio(cls, action_id=action_id, attempt=attempt),
             max_lines=max_lines,
             show_ts=show_ts,
             name=f"{action_id.run.name}:{action_id.name} ({attempt})",
+            filter_system=filter_system,
         )
         await viewer.run()

flyte/remote/_run.py

@@ -20,6 +20,8 @@ from .._protos.workflow.run_service_pb2 import WatchActionDetailsResponse
 from ._console import get_run_url
 from ._logs import Logs
 
+WaitFor = Literal["terminal", "running"]
+
 
 def _action_time_phase(action: run_definition_pb2.Action | run_definition_pb2.ActionDetails) -> rich.repr.Result:
     """
@@ -182,71 +184,28 @@ class Run:
         return run_definition_pb2.Phase.Name(self.action.phase)
 
     @syncer.wrap
-    async def wait(self, quiet: bool = False) -> None:
+    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 | None = None, max_lines: int = 100, show_ts: bool = False, raw: bool = False
+        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)
+        await self.action.show_logs(attempt, max_lines, show_ts, raw, filter_system=filter_system)
 
     async def details(self) -> RunDetails:
         """
@@ -272,15 +231,20 @@ class Run:
         )
 
     @syncer.wrap
-    async def cancel(self) -> None:
+    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:
         """
@@ -543,19 +507,25 @@ 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 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.")
+        if not details.is_running:
+            await self.wait(wait_for="running")
         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:
@@ -566,7 +536,9 @@ class Action:
             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]:
+    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.
         """
@@ -576,9 +548,80 @@ class Action:
                 return
             self._details = ad
             yield ad
+            if wait_for == "running" and ad.phase == run_definition_pb2.PHASE_RUNNING:
+                break
+            elif wait_for == "terminal" and _action_done_check(ad.phase):
+                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
+
+                    # 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.
@@ -706,6 +749,13 @@ class ActionDetails:
         """
         return run_definition_pb2.Phase.Name(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:
         """

{flyte-0.2.0b2.dist-info → flyte-0.2.0b4.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: flyte
-Version: 0.2.0b2
+Version: 0.2.0b4
 Summary: Add your description here
 Author-email: Ketan Umare <kumare3@users.noreply.github.com>
 Requires-Python: >=3.10
@@ -15,7 +15,7 @@ Requires-Dist: obstore>=0.6.0
 Requires-Dist: protobuf>=6.30.1
 Requires-Dist: pydantic>=2.10.6
 Requires-Dist: pyyaml>=6.0.2
-Requires-Dist: rich-click>=1.8.8
+Requires-Dist: rich-click>=1.8.9
 Requires-Dist: httpx>=0.28.1
 Requires-Dist: keyring>=25.6.0
 Requires-Dist: synchronicity>=0.9.11