flyte 0.2.0b3__py3-none-any.whl → 0.2.0b5__py3-none-any.whl

flyte/io/_file.py

@@ -22,7 +22,6 @@ from fsspec.asyn import AsyncFileSystem
 from fsspec.utils import get_protocol
 from mashumaro.types import SerializableType
 from pydantic import BaseModel, model_validator
-from synchronicity import Synchronizer
 
 import flyte.storage as storage
 from flyte._context import internal_ctx
@@ -33,8 +32,6 @@ from flyte.types import TypeEngine, TypeTransformer, TypeTransformerFailedError
 # Type variable for the file format
 T = TypeVar("T")
 
-synced = Synchronizer()
-
 
 class File(BaseModel, Generic[T], SerializableType):
     """
@@ -314,7 +311,7 @@ class File(BaseModel, Generic[T], SerializableType):
         with fs.open(self.path, **open_kwargs) as f:
             yield f
 
-    # @synced.wrap - enabling this did not work - synchronicity/pydantic issue
+    # TODO sync needs to be implemented
     async def download(self, local_path: Optional[Union[str, Path]] = None) -> str:
         """
         Asynchronously download the file to a local path.

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

@@ -3,25 +3,33 @@ 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
 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",
+    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,47 +57,78 @@ 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 StopAsyncIteration:
+                self.console.print("[dim]Log stream ended.[/dim]")
+            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.")
 
 
 @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).
         """
-        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
+        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(
@@ -92,6 +138,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 +148,22 @@ 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 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):
-                console.print(_format_line(line, show_ts=show_ts), end="")
+            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,
         )
         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(