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

flyte/remote/_data.py

@@ -15,7 +15,7 @@ import httpx
 from flyteidl.service import dataproxy_pb2
 from google.protobuf import duration_pb2
 
-from flyte._initialize import CommonInit, get_client, get_common_config, requires_client
+from flyte._initialize import CommonInit, ensure_client, get_client, get_common_config
 from flyte.errors import RuntimeSystemError
 
 _UPLOAD_EXPIRES_IN = timedelta(seconds=60)
@@ -109,7 +109,6 @@ async def _upload_single_file(
     return str_digest, resp.native_url
 
 
-@requires_client
 async def upload_file(fp: Path, verify: bool = True) -> Tuple[str, str]:
     """
     Uploads a file to a remote location and returns the remote URI.
@@ -119,13 +118,13 @@ async def upload_file(fp: Path, verify: bool = True) -> Tuple[str, str]:
     :return: A tuple containing the MD5 digest and the remote URI.
     """
     # This is a placeholder implementation. Replace with actual upload logic.
+    ensure_client()
     cfg = get_common_config()
     if not fp.is_file():
         raise ValueError(f"{fp} is not a single file, upload arg must be a single file.")
     return await _upload_single_file(cfg, fp, verify=verify)
 
 
-@requires_client
 async def upload_dir(dir_path: Path, verify: bool = True) -> str:
     """
     Uploads a directory to a remote location and returns the remote URI.
@@ -135,6 +134,7 @@ async def upload_dir(dir_path: Path, verify: bool = True) -> str:
     :return: The remote URI of the uploaded directory.
     """
     # This is a placeholder implementation. Replace with actual upload logic.
+    ensure_client()
     cfg = get_common_config()
     if not dir_path.is_dir():
         raise ValueError(f"{dir_path} is not a directory, upload arg must be a directory.")

flyte/remote/_logs.py

@@ -9,11 +9,11 @@ from rich.live import Live
 from rich.panel import Panel
 from rich.text import Text
 
-from flyte._api_commons import syncer
-from flyte._initialize import get_client, requires_client
+from flyte._initialize import ensure_client, get_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
+from flyte.syncify import syncify
 
 style_map = {
     payload_pb2.LogLineOriginator.SYSTEM: "bold magenta",
@@ -49,6 +49,7 @@ class AsyncLogViewer:
         name: str = "Logs",
         show_ts: bool = False,
         filter_system: bool = False,
+        panel: bool = False,
     ):
         self.console = Console()
         self.log_source = log_source
@@ -58,12 +59,15 @@ class AsyncLogViewer:
         self.show_ts = show_ts
         self.total_lines = 0
         self.filter_flyte = filter_system
+        self.panel = panel
 
-    def _render(self) -> Panel:
+    def _render(self) -> Panel | Text:
         log_text = Text()
         for line in self.lines:
             log_text.append(line)
-        return Panel(log_text, title=self.name, border_style="yellow")
+        if self.panel:
+            return Panel(log_text, title=self.name, border_style="yellow")
+        return log_text
 
     async def run(self):
         with Live(self._render(), refresh_per_second=20, console=self.console) as live:
@@ -78,6 +82,8 @@ class AsyncLogViewer:
                 pass
             except KeyboardInterrupt:
                 pass
+            except StopAsyncIteration:
+                self.console.print("[dim]Log stream ended.[/dim]")
             except LogsNotYetAvailableError as e:
                 self.console.print(f"[red]Error:[/red] {e}")
                 live.update("")
@@ -86,35 +92,47 @@ class AsyncLogViewer:
 
 @dataclass
 class Logs:
+    @syncify
     @classmethod
-    @requires_client
-    @syncer.wrap
     async def tail(
-        cls, action_id: run_definition_pb2.ActionIdentifier, attempt: int = 1
+        cls,
+        action_id: run_definition_pb2.ActionIdentifier,
+        attempt: int = 1,
+        retry: int = 3,
     ) -> AsyncGenerator[payload_pb2.LogLine, None]:
         """
         Tail the logs for a given action ID and attempt.
         :param action_id: The action ID to tail logs for.
         :param attempt: The attempt number (default is 0).
         """
-        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}."
+        ensure_client()
+        retries = 0
+        while True:
+            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
+                return
+            except asyncio.CancelledError:
+                return
+            except KeyboardInterrupt:
+                return
+            except StopAsyncIteration:
+                return
+            except grpc.aio.AioRpcError as e:
+                retries += 1
+                if retries >= retry:
+                    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}."
+                        )
+                else:
+                    await asyncio.sleep(1)
 
     @classmethod
     async def create_viewer(
@@ -125,6 +143,7 @@ class Logs:
         show_ts: bool = False,
         raw: bool = False,
         filter_system: bool = False,
+        panel: bool = False,
     ):
         """
         Create a log viewer for a given action ID and attempt.
@@ -135,19 +154,23 @@ class Logs:
         :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.
+        :param panel: Whether to use a panel for the log viewer. only applicable if raw is False.
         """
+        if attempt < 1:
+            raise ValueError("Attempt number must be greater than 0.")
         if raw:
             console = Console()
-            async for line in cls.tail.aio(cls, action_id=action_id, attempt=attempt):
+            async for line in cls.tail.aio(action_id=action_id, attempt=attempt):
                 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),
+            log_source=cls.tail.aio(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,
+            panel=panel,
         )
         await viewer.run()

flyte/remote/_project.py

@@ -1,14 +1,13 @@
 from __future__ import annotations
 
-import typing
 from dataclasses import dataclass
-from typing import AsyncGenerator, Literal, Tuple
+from typing import AsyncIterator, Iterator, Literal, Tuple, Union
 
 import rich.repr
 from flyteidl.admin import common_pb2, project_pb2
 
-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.syncify import syncify
 
 
 @dataclass
@@ -19,9 +18,8 @@ class Project:
 
     _pb2: project_pb2.Project
 
+    @syncify
     @classmethod
-    @requires_client
-    @syncer.wrap
     async def get(cls, name: str, org: str | None = None) -> Project:
         """
         Get a run by its ID or name. If both are provided, the ID will take precedence.
@@ -29,6 +27,7 @@ class Project:
         :param name: The name of the project.
         :param org: The organization of the project (if applicable).
         """
+        ensure_client()
         service = get_client().project_domain_service  # type: ignore
         resp = await service.GetProject(
             project_pb2.ProjectGetRequest(
@@ -38,14 +37,13 @@ class Project:
         )
         return cls(resp)
 
+    @syncify
     @classmethod
-    @requires_client
-    @syncer.wrap
     async def listall(
         cls,
         filters: str | None = None,
         sort_by: Tuple[str, Literal["asc", "desc"]] | None = None,
-    ) -> typing.Union[typing.Iterator[Project], AsyncGenerator[Project, None]]:
+    ) -> Union[AsyncIterator[Project], Iterator[Project]]:
         """
         Get a run by its ID or name. If both are provided, the ID will take precedence.
 
@@ -53,6 +51,7 @@ class Project:
         :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 = common_pb2.Sort(