flyte 0.2.0b1__py3-none-any.whl → 2.0.0b46__py3-none-any.whl

flyte/remote/_data.py

@@ -15,8 +15,9 @@ 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.errors import RuntimeSystemError
+from flyte._initialize import CommonInit, ensure_client, get_client, get_init_config, require_project_and_domain
+from flyte.errors import InitializationError, RuntimeSystemError
+from flyte.syncify import syncify
 
 _UPLOAD_EXPIRES_IN = timedelta(seconds=60)
 
@@ -53,9 +54,89 @@ def hash_file(file_path: typing.Union[os.PathLike, str]) -> Tuple[bytes, str, in
     return h.digest(), h.hexdigest(), size
 
 
+async def _upload_with_retry(
+    fp: Path,
+    signed_url: str,
+    extra_headers: dict,
+    verify: bool,
+    max_retries: int = 3,
+    min_backoff_sec: float = 0.5,
+    max_backoff_sec: float = 10.0,
+):
+    """
+    Upload file to signed URL with exponential backoff retry.
+
+    Retries on transient network errors and 5xx/429/408 HTTP errors.
+    Does not retry on 4xx client errors (except 408/429).
+
+    Args:
+        fp: Path to file to upload
+        signed_url: Pre-signed URL for upload
+        extra_headers: Headers including Content-MD5, Content-Length
+        verify: Whether to verify SSL certificates
+        max_retries: Maximum retry attempts (default: 3)
+        min_backoff_sec: Initial backoff delay (default: 0.5)
+        max_backoff_sec: Maximum backoff delay (default: 10.0)
+
+    Raises:
+        RuntimeSystemError: If upload fails after all retries
+    """
+    from flyte._logging import logger
+
+    retry_attempt = 0
+    last_error = None
+
+    while retry_attempt <= max_retries:
+        async with aiofiles.open(str(fp), "rb") as file:
+            async with httpx.AsyncClient(verify=verify) as aclient:
+                put_resp = await aclient.put(signed_url, headers=extra_headers, content=file)
+
+                # Success
+                if put_resp.status_code in [200, 201, 204]:
+                    if retry_attempt > 0:
+                        logger.info(f"Upload succeeded after {retry_attempt} retries for {fp.name}")
+                    return put_resp
+
+                # Check if retryable status code
+                if put_resp.status_code in [408, 429, 500, 502, 503, 504]:
+                    if retry_attempt >= max_retries:
+                        raise RuntimeSystemError(
+                            "UploadFailed",
+                            f"Failed to upload {fp} after {max_retries} retries: {put_resp.text}",
+                        )
+
+                    # Backoff and retry
+                    retry_attempt += 1
+                    if retry_attempt <= max_retries:
+                        backoff_delay = min(min_backoff_sec * (2 ** (retry_attempt - 1)), max_backoff_sec)
+                        logger.warning(
+                            f"Upload failed for {fp.name}, backing off for {backoff_delay:.2f}s "
+                            f"[retry {retry_attempt}/{max_retries}]: {last_error}"
+                        )
+                        await asyncio.sleep(backoff_delay)
+                else:
+                    # Non-retryable HTTP error
+                    raise RuntimeSystemError(
+                        "UploadFailed",
+                        f"Failed to upload {fp} to {signed_url}, status code: {put_resp.status_code}, "
+                        f"response: {put_resp.text}",
+                    )
+    return None
+
+
+@require_project_and_domain
 async def _upload_single_file(
     cfg: CommonInit, fp: Path, verify: bool = True, basedir: str | None = None
 ) -> Tuple[str, str]:
+    """
+    Upload a single file to remote storage using a signed URL.
+
+    :param cfg: Configuration containing project and domain information.
+    :param fp: Path to the file to upload.
+    :param verify: Whether to verify SSL certificates.
+    :param basedir: Optional base directory prefix for the remote path.
+    :return: Tuple of (MD5 digest hex string, remote native URL).
+    """
     md5_bytes, str_digest, _ = hash_file(fp)
     from flyte._logging import logger
 
@@ -77,39 +158,42 @@ async def _upload_single_file(
     except grpc.aio.AioRpcError as e:
         if e.code() == grpc.StatusCode.NOT_FOUND:
             raise RuntimeSystemError(
-                "NotFound", f"Failed to get signed url for {fp}, please check your project and domain."
+                "NotFound", f"Failed to get signed url for {fp}, please check your project and domain: {e.details()}"
             )
         elif e.code() == grpc.StatusCode.PERMISSION_DENIED:
             raise RuntimeSystemError(
-                "PermissionDenied", f"Failed to get signed url for {fp}, please check your permissions."
+                "PermissionDenied", f"Failed to get signed url for {fp}, please check your permissions: {e.details()}"
             )
+        elif e.code() == grpc.StatusCode.UNAVAILABLE:
+            raise InitializationError("EndpointUnavailable", "user", "Service is unavailable.")
         else:
-            raise RuntimeSystemError(e.code().value, f"Failed to get signed url for {fp}.")
+            raise RuntimeSystemError(e.code().value, f"Failed to get signed url for {fp}: {e.details()}")
     except Exception as e:
         raise RuntimeSystemError(type(e).__name__, f"Failed to get signed url for {fp}.") from e
-    logger.debug(f"Uploading to signed url {resp.signed_url} for {fp}")
+    logger.debug(f"Uploading to [link={resp.signed_url}]signed url[/link] for [link=file://{fp}]{fp}[/link]")
     extra_headers = get_extra_headers_for_protocol(resp.native_url)
     extra_headers.update(resp.headers)
     encoded_md5 = b64encode(md5_bytes)
     content_length = fp.stat().st_size
 
-    async with aiofiles.open(str(fp), "rb") as file:
-        extra_headers.update({"Content-Length": str(content_length), "Content-MD5": encoded_md5.decode("utf-8")})
-        async with httpx.AsyncClient(verify=verify) as aclient:
-            put_resp = await aclient.put(resp.signed_url, headers=extra_headers, content=file)
-            if put_resp.status_code != 200:
-                raise RuntimeSystemError(
-                    "UploadFailed",
-                    f"Failed to upload {fp} to {resp.signed_url}, status code: {put_resp.status_code}",
-                )
-        # TODO in old code we did this
-        #             if self._config.platform.insecure_skip_verify is True
-        #             else self._config.platform.ca_cert_file_path,
+    # Update headers with MD5 and content length
+    extra_headers.update({"Content-Length": str(content_length), "Content-MD5": encoded_md5.decode("utf-8")})
+
+    await _upload_with_retry(
+        fp=fp,
+        signed_url=resp.signed_url,
+        extra_headers=extra_headers,
+        verify=verify,
+        max_retries=3,
+        min_backoff_sec=0.5,
+        max_backoff_sec=10.0,
+    )
+
     logger.debug(f"Uploaded with digest {str_digest}, blob location is {resp.native_url}")
     return str_digest, resp.native_url
 
 
-@requires_client
+@syncify
 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 +203,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.
-    cfg = get_common_config()
+    ensure_client()
+    cfg = get_init_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,7 +219,8 @@ 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.
-    cfg = get_common_config()
+    ensure_client()
+    cfg = get_init_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,44 @@ from collections import deque
 from dataclasses import dataclass
 from typing import AsyncGenerator, AsyncIterator
 
+import grpc
+from flyteidl2.common import identifier_pb2
+from flyteidl2.logs.dataplane import payload_pb2
+from flyteidl2.workflow import run_logs_service_pb2
 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._protos.logs.dataplane import payload_pb2
-from flyte._protos.workflow import run_definition_pb2, run_logs_service_pb2
+from flyte._initialize import ensure_client, get_client
+from flyte._logging import logger
+from flyte._tools import ipython_check, ipywidgets_check
+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:
+    """
+    Format a log line for display with optional timestamp and system filtering.
+
+    :param logline: The log line protobuf to format.
+    :param show_ts: Whether to include timestamps.
+    :param filter_system: Whether to filter out system log lines.
+    :return: A formatted Text object or None if the line should be filtered out.
+    """
+    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 "[flyte]" in logline.message and "flyte.errors" not in logline.message:
+        if filter_system:
+            return None
         style = "dim"
     ts = ""
     if show_ts:
@@ -34,7 +53,15 @@ 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,
+        panel: bool = False,
+    ):
         self.console = Console()
         self.log_source = log_source
         self.max_lines = max_lines
@@ -42,56 +69,98 @@ class AsyncLogViewer:
         self.name = name
         self.show_ts = show_ts
         self.total_lines = 0
+        self.filter_flyte = filter_system
+        self.panel = panel
 
-    def _render(self):
+    def _render(self) -> Panel | Text:
+        """
+        Render the current log lines as a Panel or Text object for display.
+        """
         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=10, console=self.console) as live:
+        """
+        Run the log viewer, streaming and displaying log lines until completion.
+        """
+        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: identifier_pb2.ActionIdentifier,
+        attempt: int = 1,
+        retry: int = 5,
     ) -> 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(2)
 
     @classmethod
     async def create_viewer(
         cls,
-        action_id: run_definition_pb2.ActionIdentifier,
+        action_id: identifier_pb2.ActionIdentifier,
         attempt: int = 1,
         max_lines: int = 30,
         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.
@@ -101,16 +170,30 @@ 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.
+        :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 ipython_check():
+            if not ipywidgets_check():
+                logger.warning("IPython widgets is not available, defaulting to console output.")
+                raw = True
+
         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}",
+            name=f"{action_id.run.name}:{action_id.name} ({attempt})",
+            filter_system=filter_system,
+            panel=panel,
         )
         await viewer.run()

flyte/remote/_project.py

@@ -1,27 +1,28 @@
 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
+from flyte.syncify import syncify
 
+from ._common import ToJSONMixin
 
+
+# TODO Add support for orgs again
 @dataclass
-class Project:
+class Project(ToJSONMixin):
     """
     A class representing a project in the Union API.
     """
 
-    _pb2: project_pb2.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,23 +30,23 @@ 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(
                 id=name,
-                org=org,
+                # org=org,
             )
         )
         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,12 +54,13 @@ 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(
             key=sort_by[0], direction=common_pb2.Sort.ASCENDING if sort_by[1] == "asc" else common_pb2.Sort.DESCENDING
         )
-        org = get_common_config().org
+        # org = get_common_config().org
         while True:
             resp = await get_client().project_domain_service.ListProjects(  # type: ignore
                 project_pb2.ProjectListRequest(
@@ -66,7 +68,7 @@ class Project:
                     token=token,
                     filters=filters,
                     sort_by=sort_pb2,
-                    org=org,
+                    # org=org,
                 )
             )
             token = resp.token
@@ -76,11 +78,11 @@ class Project:
                 break
 
     def __rich_repr__(self) -> rich.repr.Result:
-        yield "name", self._pb2.name
-        yield "id", self._pb2.id
-        yield "description", self._pb2.description
-        yield "state", project_pb2.Project.ProjectState.Name(self._pb2.state)
+        yield "name", self.pb2.name
+        yield "id", self.pb2.id
+        yield "description", self.pb2.description
+        yield "state", project_pb2.Project.ProjectState.Name(self.pb2.state)
         yield (
             "labels",
-            ", ".join([f"{k}: {v}" for k, v in self._pb2.labels.values.items()]) if self._pb2.labels else None,
+            ", ".join([f"{k}: {v}" for k, v in self.pb2.labels.values.items()]) if self.pb2.labels else None,
         )